Skip to main content

mentedb_server/
handlers_memory.rs

1//! Memory feature endpoints: hybrid search, heads-up context, history, graph, conflicts, pins, profile, entities, batch store.
2use std::collections::{HashMap, HashSet, VecDeque};
3use std::sync::Arc;
4
5use axum::extract::{Path, Query};
6use axum::http::StatusCode;
7use axum::response::IntoResponse;
8use axum::routing::{get, post};
9use axum::{Extension, Json, Router};
10use mentedb::injection::{InjectionQuery, SelectionReason};
11use mentedb_core::edge::EdgeType;
12use mentedb_core::error::MenteError;
13use mentedb_core::types::{AgentId, MemoryId, SpaceId, UserId};
14use mentedb_core::{MemoryEdge, MemoryNode};
15use serde_json::{Value, json};
16use tracing::error;
17use uuid::Uuid;
18
19use crate::auth::AuthenticatedAgent;
20use crate::databases::Db;
21use crate::error::ApiError;
22use crate::handlers::{
23    check_body_agent, check_owner, embed_query, memory_id_param, memory_node_to_json,
24    parse_attributes, parse_edge_type, parse_embedding, parse_memory_type, parse_optional_uuid,
25    parse_uuid,
26};
27use crate::state::AppState;
28
29/// Tag the injection path treats as pinned: always injected, never ranked out.
30const PIN_TAG: &str = "scope:always";
31/// Tag marking a pin made by hand, which standing-rule pruning keeps.
32const MANUAL_TAG: &str = "source:manual";
33/// Tag prefix of entity memories.
34const ENTITY_TAG_PREFIX: &str = "entity:";
35
36/// Request size and result size limits for these endpoints.
37#[derive(Clone, Debug)]
38pub struct MemoryApiLimits {
39    /// Results returned by hybrid search when `k` is not given.
40    pub default_k: usize,
41    /// Most results one hybrid search may ask for.
42    pub max_k: usize,
43    /// Hops walked by the neighbors endpoint when `depth` is not given.
44    pub default_neighbor_depth: usize,
45    /// Most hops the neighbors endpoint walks.
46    pub max_neighbor_depth: usize,
47    /// Most memories the neighbors endpoint returns.
48    pub max_neighbors: usize,
49    /// Most predecessors the history endpoint walks.
50    pub max_history_depth: usize,
51    /// Items returned by list endpoints when `limit` is not given.
52    pub default_list: usize,
53    /// Most items one list request may ask for.
54    pub max_list: usize,
55    /// Most memories one batch store may carry.
56    pub max_batch: usize,
57    /// Ranked context items (beyond pins) when `max_items` is not given.
58    pub default_context_items: usize,
59    /// Most ranked context items one request may ask for.
60    pub max_context_items: usize,
61    /// Verbatim episodic items within the context when not given.
62    pub default_context_episodic: usize,
63}
64
65impl Default for MemoryApiLimits {
66    fn default() -> Self {
67        Self {
68            default_k: 10,
69            max_k: 200,
70            default_neighbor_depth: 1,
71            max_neighbor_depth: 4,
72            max_neighbors: 500,
73            max_history_depth: 64,
74            default_list: 50,
75            max_list: 1000,
76            max_batch: 1000,
77            default_context_items: 10,
78            max_context_items: 100,
79            default_context_episodic: 3,
80        }
81    }
82}
83
84type Limits = Extension<Arc<MemoryApiLimits>>;
85type Caller = Option<Extension<AuthenticatedAgent>>;
86
87/// Routes on one database, mounted under `/v1` and `/v1/db/{db}`.
88pub fn routes() -> Router<Arc<AppState>> {
89    routes_with_limits(MemoryApiLimits::default())
90}
91
92/// [`routes`] with custom limits.
93pub fn routes_with_limits(limits: MemoryApiLimits) -> Router<Arc<AppState>> {
94    Router::new()
95        .route("/search/hybrid", post(search_hybrid))
96        .route("/context", post(context))
97        .route("/memories/batch", post(store_batch))
98        .route("/memories/{id}/history", get(history))
99        .route("/memories/{id}/neighbors", get(neighbors))
100        .route("/memories/{id}/pin", post(pin).delete(unpin))
101        .route("/pins", get(list_pins))
102        .route("/conflicts", get(conflicts))
103        .route("/profile", get(profile))
104        .route("/entities", get(entities))
105        .route("/entities/{name}/memories", get(entity_memories))
106        .route("/recent", get(recent))
107        .layer(Extension(Arc::new(limits)))
108}
109
110// ---------------------------------------------------------------------------
111// POST /v1/search/hybrid
112// ---------------------------------------------------------------------------
113
114/// Hybrid recall: vector similarity fused with BM25 keyword matching, at a
115/// point in time, optionally scoped to an agent and a user.
116async fn search_hybrid(
117    agent: Caller,
118    Extension(limits): Limits,
119    db: Db,
120    Json(req): Json<Value>,
121) -> Result<impl IntoResponse, ApiError> {
122    let k = bounded(body_usize(&req, "k")?, limits.default_k, limits.max_k);
123    let query_text = body_str(&req, "query");
124    let embedding = query_embedding(&db, &req, query_text.as_deref()).await?;
125    let at = match req.get("as_of") {
126        Some(v) => v
127            .as_u64()
128            .ok_or_else(|| ApiError::BadRequest("'as_of' must be microseconds".into()))?,
129        None => now_micros(),
130    };
131    let time_range = match req.get("time_range") {
132        Some(v) => {
133            let pair = v
134                .as_array()
135                .filter(|a| a.len() == 2)
136                .and_then(|a| Some((a[0].as_u64()?, a[1].as_u64()?)))
137                .ok_or_else(|| {
138                    ApiError::BadRequest(
139                        "'time_range' must be [from, until] in microseconds".into(),
140                    )
141                })?;
142            Some(pair)
143        }
144        None => None,
145    };
146    let tags = string_array(&req, "tags")?;
147    let tags_or = match req.get("tags_mode").and_then(|v| v.as_str()) {
148        None | Some("and") => false,
149        Some("or") => true,
150        Some(other) => {
151            return Err(ApiError::BadRequest(format!(
152                "unknown tags_mode '{other}': use 'and' or 'or'"
153            )));
154        }
155    };
156    let agent_scope = scope_agent(&agent, parse_optional_uuid(&req, "agent_id")?)?;
157    let user_scope = parse_optional_uuid(&req, "user_id")?.map(UserId);
158    let project = body_str(&req, "project");
159
160    let handle = db.handle.clone();
161    let results = tokio::task::spawn_blocking(move || {
162        let tag_refs: Vec<&str> = tags.iter().map(String::as_str).collect();
163        let hits = handle.recall_hybrid_scoped_at_mode(
164            &embedding,
165            query_text.as_deref(),
166            k,
167            at,
168            (!tag_refs.is_empty()).then_some(tag_refs.as_slice()),
169            tags_or,
170            time_range,
171            agent_scope,
172            user_scope,
173            project.as_deref(),
174        )?;
175        Ok::<_, MenteError>(
176            hits.into_iter()
177                .filter_map(|(id, score)| {
178                    let node = handle.get_memory(id).ok()?;
179                    Some(with_fields(&node, json!({ "score": score })))
180                })
181                .collect::<Vec<Value>>(),
182        )
183    })
184    .await
185    .map_err(task_failed)?
186    .map_err(|e| engine_error("hybrid search", e))?;
187
188    Ok(Json(json!({
189        "database": db.name,
190        "as_of": at,
191        "results": results,
192    })))
193}
194
195// ---------------------------------------------------------------------------
196// POST /v1/context
197// ---------------------------------------------------------------------------
198
199/// Heads-up context for a prompt: pinned memories plus the relevant ones,
200/// each with its score and the reason it was selected.
201async fn context(
202    agent: Caller,
203    Extension(limits): Limits,
204    db: Db,
205    Json(req): Json<Value>,
206) -> Result<impl IntoResponse, ApiError> {
207    let query_text = body_str(&req, "query");
208    let embedding = query_embedding(&db, &req, query_text.as_deref()).await?;
209    let max_items = bounded(
210        body_usize(&req, "max_items")?,
211        limits.default_context_items,
212        limits.max_context_items,
213    );
214    let max_episodic = body_usize(&req, "max_episodic")?
215        .unwrap_or(limits.default_context_episodic)
216        .min(max_items);
217    let session_id = body_str(&req, "session_id");
218    let project = body_str(&req, "project");
219    let exclude_ids = string_array(&req, "exclude_ids")?
220        .iter()
221        .map(|s| {
222            s.parse::<MemoryId>()
223                .map_err(|_| ApiError::BadRequest(format!("invalid id in 'exclude_ids': {s}")))
224        })
225        .collect::<Result<Vec<_>, _>>()?;
226    let agent_scope = scope_agent(&agent, parse_optional_uuid(&req, "agent_id")?)?;
227    let user_scope = parse_optional_uuid(&req, "user_id")?.map(UserId);
228
229    let handle = db.handle.clone();
230    let selected = tokio::task::spawn_blocking(move || {
231        let query = InjectionQuery {
232            embedding: &embedding,
233            query_text: query_text.as_deref(),
234            session_id: session_id.as_deref(),
235            exclude_ids: &exclude_ids,
236            max_items,
237            max_episodic,
238            agent_id: agent_scope,
239            user_id: user_scope,
240            current_project: project.as_deref(),
241        };
242        handle.recall_for_injection(&query)
243    })
244    .await
245    .map_err(task_failed)?
246    .map_err(|e| engine_error("context recall", e))?;
247
248    let memories: Vec<Value> = selected
249        .iter()
250        .map(|c| {
251            let reason = match c.reason {
252                SelectionReason::Pinned => "pinned",
253                SelectionReason::Relevant => "relevant",
254                SelectionReason::ModeRule => "mode_rule",
255            };
256            with_fields(&c.node, json!({ "score": c.score, "reason": reason }))
257        })
258        .collect();
259    Ok(Json(json!({ "database": db.name, "memories": memories })))
260}
261
262// ---------------------------------------------------------------------------
263// GET /v1/memories/{id}/history
264// ---------------------------------------------------------------------------
265
266/// The chain of older facts this memory replaced, newest predecessor first.
267async fn history(
268    agent: Caller,
269    Extension(limits): Limits,
270    db: Db,
271    Path(params): Path<HashMap<String, String>>,
272    Query(query): Query<HashMap<String, String>>,
273) -> Result<impl IntoResponse, ApiError> {
274    let id = memory_id_param(&params)?;
275    let depth = bounded(
276        query_usize(&query, "depth")?,
277        limits.max_history_depth,
278        limits.max_history_depth,
279    );
280    let node = load_owned(&agent, &db, id)?;
281    let handle = db.handle.clone();
282    let chain = tokio::task::spawn_blocking(move || handle.supersession_history(id, depth))
283        .await
284        .map_err(task_failed)?
285        .map_err(|e| engine_error("history", e))?;
286    Ok(Json(json!({
287        "memory": memory_node_to_json(&node),
288        "history": chain.iter().map(memory_node_to_json).collect::<Vec<_>>(),
289    })))
290}
291
292// ---------------------------------------------------------------------------
293// GET /v1/memories/{id}/neighbors
294// ---------------------------------------------------------------------------
295
296/// Memories within `depth` hops of one memory in the graph, with the edges
297/// between them. `edge_types` (comma separated) limits which edges are walked;
298/// `direction` is `out`, `in` or `both` (the default).
299async fn neighbors(
300    agent: Caller,
301    Extension(limits): Limits,
302    db: Db,
303    Path(params): Path<HashMap<String, String>>,
304    Query(query): Query<HashMap<String, String>>,
305) -> Result<impl IntoResponse, ApiError> {
306    let id = memory_id_param(&params)?;
307    let depth = bounded(
308        query_usize(&query, "depth")?,
309        limits.default_neighbor_depth,
310        limits.max_neighbor_depth,
311    );
312    let edge_types: Vec<EdgeType> = match query.get("edge_types").filter(|s| !s.is_empty()) {
313        Some(list) => list
314            .split(',')
315            .map(|s| {
316                parse_edge_type(s.trim())
317                    .ok_or_else(|| ApiError::BadRequest(format!("unknown edge_type: {s}")))
318            })
319            .collect::<Result<_, _>>()?,
320        None => Vec::new(),
321    };
322    let (outgoing, incoming) = match query.get("direction").map(String::as_str) {
323        None | Some("both") => (true, true),
324        Some("out") => (true, false),
325        Some("in") => (false, true),
326        Some(other) => {
327            return Err(ApiError::BadRequest(format!(
328                "unknown direction '{other}': use 'out', 'in' or 'both'"
329            )));
330        }
331    };
332    let max_nodes = limits.max_neighbors;
333    let owner = token_agent(&agent)?;
334    let center = load_owned(&agent, &db, id)?;
335
336    let handle = db.handle.clone();
337    let (nodes, edges) = tokio::task::spawn_blocking(move || {
338        let wanted = |t: EdgeType| edge_types.is_empty() || edge_types.contains(&t);
339        // Walk under the graph read lock, then release it before loading
340        // memories: the engine's lock order takes the page map first.
341        let (reached, edges) = {
342            let graph = handle.graph().graph();
343            let mut seen: HashSet<MemoryId> = HashSet::from([id]);
344            let mut reached: Vec<(MemoryId, usize)> = Vec::new();
345            let mut edges: Vec<MemoryEdge> = Vec::new();
346            let mut queue = VecDeque::from([(id, 0usize)]);
347            while let Some((node, hops)) = queue.pop_front() {
348                if hops >= depth {
349                    continue;
350                }
351                let mut steps = Vec::new();
352                if outgoing {
353                    for (to, e) in graph.outgoing(node) {
354                        steps.push((to, stored_edge(node, to, &e)));
355                    }
356                }
357                if incoming {
358                    for (from, e) in graph.incoming(node) {
359                        steps.push((from, stored_edge(from, node, &e)));
360                    }
361                }
362                for (next, edge) in steps {
363                    if !wanted(edge.edge_type) {
364                        continue;
365                    }
366                    if seen.contains(&next) {
367                        edges.push(edge);
368                        continue;
369                    }
370                    if reached.len() >= max_nodes {
371                        continue;
372                    }
373                    seen.insert(next);
374                    reached.push((next, hops + 1));
375                    edges.push(edge);
376                    queue.push_back((next, hops + 1));
377                }
378            }
379            (reached, edges)
380        };
381        let mut dedup = HashSet::new();
382        let edges: Vec<MemoryEdge> = edges
383            .into_iter()
384            .filter(|e| dedup.insert((e.source, e.target, e.edge_type)))
385            .collect();
386        let nodes: Vec<(MemoryNode, usize)> = reached
387            .into_iter()
388            .filter_map(|(nid, hops)| handle.get_memory(nid).ok().map(|n| (n, hops)))
389            .filter(|(n, _)| visible(owner, n))
390            .collect();
391        (nodes, edges)
392    })
393    .await
394    .map_err(task_failed)?;
395
396    let shown: HashSet<MemoryId> = nodes
397        .iter()
398        .map(|(n, _)| n.id)
399        .chain(std::iter::once(id))
400        .collect();
401    Ok(Json(json!({
402        "memory": memory_node_to_json(&center),
403        "depth": depth,
404        "neighbors": nodes
405            .iter()
406            .map(|(n, hops)| with_fields(n, json!({ "hops": hops })))
407            .collect::<Vec<_>>(),
408        "edges": edges
409            .iter()
410            .filter(|e| shown.contains(&e.source) && shown.contains(&e.target))
411            .map(edge_to_json)
412            .collect::<Vec<_>>(),
413    })))
414}
415
416// ---------------------------------------------------------------------------
417// GET /v1/conflicts
418// ---------------------------------------------------------------------------
419
420/// Contradiction and supersession edges, with both memories of each pair.
421async fn conflicts(
422    agent: Caller,
423    Extension(limits): Limits,
424    db: Db,
425    Query(query): Query<HashMap<String, String>>,
426) -> Result<impl IntoResponse, ApiError> {
427    let limit = list_limit(&query, &limits)?;
428    let owner = token_agent(&agent)?;
429    let handle = db.handle.clone();
430    let (counts, pairs) = tokio::task::spawn_blocking(move || {
431        let counts = handle.conflict_edge_counts();
432        let mut seen = HashSet::new();
433        let mut pairs = Vec::new();
434        for (source, target, edge_type, weight) in handle.conflict_edges() {
435            if pairs.len() >= limit {
436                break;
437            }
438            // A contradiction is undirected; show each pair once.
439            let key = if edge_type == EdgeType::Contradicts && target < source {
440                (target, source, edge_type)
441            } else {
442                (source, target, edge_type)
443            };
444            if !seen.insert(key) {
445                continue;
446            }
447            let (Ok(a), Ok(b)) = (handle.get_memory(source), handle.get_memory(target)) else {
448                continue;
449            };
450            if visible(owner, &a) && visible(owner, &b) {
451                pairs.push(json!({
452                    "edge_type": edge_type_name(edge_type),
453                    "weight": weight,
454                    "source": memory_node_to_json(&a),
455                    "target": memory_node_to_json(&b),
456                }));
457            }
458        }
459        (counts, pairs)
460    })
461    .await
462    .map_err(task_failed)?;
463    Ok(Json(json!({
464        "contradicts": counts.0,
465        "supersedes": counts.1,
466        "conflicts": pairs,
467    })))
468}
469
470// ---------------------------------------------------------------------------
471// Pins: GET /v1/pins, POST and DELETE /v1/memories/{id}/pin
472// ---------------------------------------------------------------------------
473
474/// Pinned memories, which every heads-up context includes.
475async fn list_pins(
476    agent: Caller,
477    Extension(limits): Limits,
478    db: Db,
479    Query(query): Query<HashMap<String, String>>,
480) -> Result<impl IntoResponse, ApiError> {
481    let limit = list_limit(&query, &limits)?;
482    let owner = token_agent(&agent)?;
483    let handle = db.handle.clone();
484    let pins = tokio::task::spawn_blocking(move || {
485        let ids = handle.memory_ids_with_tag(PIN_TAG);
486        newest_visible(&handle, ids, owner, limit)
487    })
488    .await
489    .map_err(task_failed)?;
490    Ok(Json(json!({ "pins": pins })))
491}
492
493/// Pin a memory so every heads-up context includes it.
494async fn pin(
495    agent: Caller,
496    db: Db,
497    Path(params): Path<HashMap<String, String>>,
498) -> Result<impl IntoResponse, ApiError> {
499    set_pin(agent, db, params, true).await
500}
501
502/// Remove a memory's pin. The memory itself is kept.
503async fn unpin(
504    agent: Caller,
505    db: Db,
506    Path(params): Path<HashMap<String, String>>,
507) -> Result<impl IntoResponse, ApiError> {
508    set_pin(agent, db, params, false).await
509}
510
511async fn set_pin(
512    agent: Caller,
513    db: Db,
514    params: HashMap<String, String>,
515    pinned: bool,
516) -> Result<Json<Value>, ApiError> {
517    let id = memory_id_param(&params)?;
518    let mut node = load_owned(&agent, &db, id)?;
519    if node.tags.iter().any(|t| t == PIN_TAG) != pinned {
520        if pinned {
521            node.tags.push(PIN_TAG.to_string());
522            if !node.tags.iter().any(|t| t == MANUAL_TAG) {
523                node.tags.push(MANUAL_TAG.to_string());
524            }
525        } else {
526            node.tags.retain(|t| t != PIN_TAG);
527        }
528        // Re-storing under the same id replaces the node and re-indexes its
529        // tags, the engine's own path for pinning and un-pinning.
530        let handle = db.handle.clone();
531        tokio::task::spawn_blocking(move || handle.store(node))
532            .await
533            .map_err(task_failed)?
534            .map_err(|e| engine_error("pin", e))?;
535    }
536    let now = db
537        .get_memory(id)
538        .map_err(|_| ApiError::NotFound(format!("memory {id} not found")))?;
539    let is_pinned = now.tags.iter().any(|t| t == PIN_TAG);
540    if is_pinned != pinned {
541        // The engine skips a re-store that duplicates another memory of the
542        // same owner, such as an identical standing rule already pinned.
543        return Err(ApiError::Conflict(format!(
544            "memory {id} duplicates another memory of the same owner; its pin was not changed"
545        )));
546    }
547    Ok(Json(with_fields(&now, json!({ "pinned": is_pinned }))))
548}
549
550// ---------------------------------------------------------------------------
551// GET /v1/profile
552// ---------------------------------------------------------------------------
553
554/// The user profile and the facts it is built from, for one owner pair.
555async fn profile(
556    agent: Caller,
557    db: Db,
558    Query(query): Query<HashMap<String, String>>,
559) -> Result<impl IntoResponse, ApiError> {
560    let user = query_uuid(&query, "user_id")?
561        .map(UserId)
562        .unwrap_or_else(UserId::nil);
563    let owner = scope_agent(&agent, query_uuid(&query, "agent_id")?)?.unwrap_or_else(AgentId::nil);
564    let handle = db.handle.clone();
565    let (profile, facts) = tokio::task::spawn_blocking(move || {
566        (
567            handle.user_profile_for(user, owner),
568            handle.profile_facts(user, owner),
569        )
570    })
571    .await
572    .map_err(task_failed)?;
573    Ok(Json(json!({
574        "user_id": user.to_string(),
575        "agent_id": owner.to_string(),
576        "profile": profile.as_ref().map(memory_node_to_json),
577        "facts": facts,
578    })))
579}
580
581// ---------------------------------------------------------------------------
582// GET /v1/entities, GET /v1/entities/{name}/memories
583// ---------------------------------------------------------------------------
584
585/// Entities named by entity memories, with how many memories carry each,
586/// plus the names the entity resolver knows.
587async fn entities(
588    agent: Caller,
589    Extension(limits): Limits,
590    db: Db,
591    Query(query): Query<HashMap<String, String>>,
592) -> Result<impl IntoResponse, ApiError> {
593    let limit = list_limit(&query, &limits)?;
594    let owner = token_agent(&agent)?;
595    let handle = db.handle.clone();
596    let (counts, known) = tokio::task::spawn_blocking(move || {
597        let mut counts: HashMap<String, usize> = HashMap::new();
598        for node in handle.entity_memories() {
599            if !visible(owner, &node) {
600                continue;
601            }
602            for tag in &node.tags {
603                if let Some(name) = tag.strip_prefix(ENTITY_TAG_PREFIX) {
604                    *counts.entry(name.to_string()).or_default() += 1;
605                }
606            }
607        }
608        (counts, handle.known_entities())
609    })
610    .await
611    .map_err(task_failed)?;
612    let mut names: Vec<(String, usize)> = counts.into_iter().collect();
613    names.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
614    names.truncate(limit);
615    Ok(Json(json!({
616        "entities": names
617            .iter()
618            .map(|(name, n)| json!({ "name": name, "memory_count": n }))
619            .collect::<Vec<_>>(),
620        // The resolver is shared by the whole database, so a caller bound to
621        // one agent does not see it.
622        "resolver": if owner.is_none() { known } else { Vec::new() },
623    })))
624}
625
626/// Memories tagged with an entity, after resolving the name to its canonical
627/// form.
628async fn entity_memories(
629    agent: Caller,
630    Extension(limits): Limits,
631    db: Db,
632    Path(params): Path<HashMap<String, String>>,
633    Query(query): Query<HashMap<String, String>>,
634) -> Result<impl IntoResponse, ApiError> {
635    let name = params
636        .get("name")
637        .map(|s| s.trim().to_string())
638        .filter(|s| !s.is_empty())
639        .ok_or_else(|| ApiError::BadRequest("missing entity name".into()))?;
640    let limit = list_limit(&query, &limits)?;
641    let owner = token_agent(&agent)?;
642    let resolved = db.resolve_entity(&name);
643    let mut tags = vec![format!("{ENTITY_TAG_PREFIX}{}", name.to_lowercase())];
644    let canonical_tag = format!(
645        "{ENTITY_TAG_PREFIX}{}",
646        resolved.canonical.to_lowercase().trim()
647    );
648    if !tags.contains(&canonical_tag) {
649        tags.push(canonical_tag);
650    }
651    let handle = db.handle.clone();
652    let memories = tokio::task::spawn_blocking(move || {
653        let mut ids: Vec<MemoryId> = Vec::new();
654        for tag in &tags {
655            for id in handle.memory_ids_with_tag(tag) {
656                if !ids.contains(&id) {
657                    ids.push(id);
658                }
659            }
660        }
661        newest_visible(&handle, ids, owner, limit)
662    })
663    .await
664    .map_err(task_failed)?;
665    Ok(Json(json!({
666        "name": name,
667        "canonical": resolved.canonical,
668        "confidence": resolved.confidence,
669        "source": format!("{:?}", resolved.source).to_lowercase(),
670        "memories": memories,
671    })))
672}
673
674// ---------------------------------------------------------------------------
675// POST /v1/memories/batch
676// ---------------------------------------------------------------------------
677
678/// Store many memories in one call. Every item is validated before any is
679/// stored. Items without an `embedding` are embedded by the server's embedder
680/// when one is configured.
681async fn store_batch(
682    agent: Caller,
683    Extension(limits): Limits,
684    db: Db,
685    Json(req): Json<Value>,
686) -> Result<impl IntoResponse, ApiError> {
687    let items = req
688        .get("memories")
689        .and_then(|v| v.as_array())
690        .ok_or_else(|| ApiError::BadRequest("missing 'memories' array".into()))?;
691    if items.is_empty() {
692        return Err(ApiError::BadRequest("'memories' is empty".into()));
693    }
694    if items.len() > limits.max_batch {
695        return Err(ApiError::BadRequest(format!(
696            "batch of {} exceeds the limit of {}",
697            items.len(),
698            limits.max_batch
699        )));
700    }
701    let now = now_micros();
702    let nodes = items
703        .iter()
704        .enumerate()
705        .map(|(i, item)| parse_node(&agent, item, now).map_err(|e| at_index(i, e)))
706        .collect::<Result<Vec<_>, _>>()?;
707
708    let handle = db.handle.clone();
709    let ids = tokio::task::spawn_blocking(move || {
710        let mut nodes = nodes;
711        let missing: Vec<usize> = nodes
712            .iter()
713            .enumerate()
714            .filter(|(_, n)| n.embedding.is_empty())
715            .map(|(i, _)| i)
716            .collect();
717        if !missing.is_empty() {
718            let texts: Vec<&str> = missing.iter().map(|&i| nodes[i].content.as_str()).collect();
719            if let Some(vectors) = handle.embed_batch(&texts)? {
720                for (i, v) in missing.into_iter().zip(vectors) {
721                    nodes[i].embedding = v;
722                }
723            }
724        }
725        handle.store_batch(nodes)
726    })
727    .await
728    .map_err(task_failed)?
729    .map_err(|e| match e {
730        MenteError::EmbeddingDimensionMismatch { .. } => ApiError::BadRequest(e.to_string()),
731        e => engine_error("batch store", e),
732    })?;
733
734    Ok((
735        StatusCode::CREATED,
736        Json(json!({
737            "ids": ids.iter().map(|id| id.to_string()).collect::<Vec<_>>(),
738            "stored": ids.len(),
739        })),
740    ))
741}
742
743/// One batch item as a memory node, with the same fields as `POST /memories`.
744fn parse_node(agent: &Caller, req: &Value, now: u64) -> Result<MemoryNode, ApiError> {
745    let agent_id = AgentId(parse_uuid(req, "agent_id")?);
746    check_body_agent(agent, agent_id)?;
747    let memory_type = parse_memory_type(req)?;
748    let content = req
749        .get("content")
750        .and_then(|v| v.as_str())
751        .ok_or_else(|| ApiError::BadRequest("missing or invalid 'content'".into()))?
752        .to_string();
753    let space_id = match req.get("space_id") {
754        Some(v) => v
755            .as_str()
756            .and_then(|s| s.parse::<SpaceId>().ok())
757            .ok_or_else(|| ApiError::BadRequest("invalid 'space_id'".into()))?,
758        None => SpaceId::nil(),
759    };
760    let optional_u64 = |field: &str| match req.get(field) {
761        None | Some(Value::Null) => Ok(None),
762        Some(v) => v
763            .as_u64()
764            .map(Some)
765            .ok_or_else(|| ApiError::BadRequest(format!("'{field}' must be microseconds"))),
766    };
767    Ok(MemoryNode {
768        id: MemoryId::new(),
769        agent_id,
770        user_id: parse_optional_uuid(req, "user_id")?
771            .map(UserId)
772            .unwrap_or_else(UserId::nil),
773        memory_type,
774        embedding: parse_embedding(req, false)?,
775        content,
776        created_at: now,
777        accessed_at: now,
778        access_count: 0,
779        salience: req.get("salience").and_then(|v| v.as_f64()).unwrap_or(0.5) as f32,
780        confidence: req
781            .get("confidence")
782            .and_then(|v| v.as_f64())
783            .unwrap_or(1.0) as f32,
784        space_id,
785        attributes: parse_attributes(req.get("attributes")),
786        tags: string_array(req, "tags")?,
787        valid_from: optional_u64("valid_from")?,
788        valid_until: optional_u64("valid_until")?,
789        context: None,
790    })
791}
792
793fn at_index(i: usize, e: ApiError) -> ApiError {
794    match e {
795        ApiError::BadRequest(m) => ApiError::BadRequest(format!("memories[{i}]: {m}")),
796        ApiError::Forbidden(m) => ApiError::Forbidden(format!("memories[{i}]: {m}")),
797        other => other,
798    }
799}
800
801// ---------------------------------------------------------------------------
802// GET /v1/recent
803// ---------------------------------------------------------------------------
804
805/// The newest memories, newest first.
806async fn recent(
807    agent: Caller,
808    Extension(limits): Limits,
809    db: Db,
810    Query(query): Query<HashMap<String, String>>,
811) -> Result<impl IntoResponse, ApiError> {
812    let limit = list_limit(&query, &limits)?;
813    let owner = token_agent(&agent)?;
814    let handle = db.handle.clone();
815    let memories = tokio::task::spawn_blocking(move || {
816        handle
817            .recall_recent(limit)
818            .iter()
819            .filter(|n| visible(owner, n))
820            .map(memory_node_to_json)
821            .collect::<Vec<_>>()
822    })
823    .await
824    .map_err(task_failed)?;
825    Ok(Json(json!({ "memories": memories })))
826}
827
828// ---------------------------------------------------------------------------
829// Helpers
830// ---------------------------------------------------------------------------
831
832fn now_micros() -> u64 {
833    std::time::SystemTime::now()
834        .duration_since(std::time::UNIX_EPOCH)
835        .unwrap_or_default()
836        .as_micros() as u64
837}
838
839fn task_failed(e: tokio::task::JoinError) -> ApiError {
840    ApiError::Internal(format!("engine task failed: {e}"))
841}
842
843fn engine_error(what: &str, e: MenteError) -> ApiError {
844    error!("{what} failed: {e}");
845    ApiError::Internal(format!("{what} failed: {e}"))
846}
847
848/// `value` when given (at least 1), else `default`, never above `max`.
849fn bounded(value: Option<usize>, default: usize, max: usize) -> usize {
850    value.unwrap_or(default).clamp(1, max.max(1))
851}
852
853fn list_limit(
854    query: &HashMap<String, String>,
855    limits: &MemoryApiLimits,
856) -> Result<usize, ApiError> {
857    Ok(bounded(
858        query_usize(query, "limit")?,
859        limits.default_list,
860        limits.max_list,
861    ))
862}
863
864fn query_usize(query: &HashMap<String, String>, field: &str) -> Result<Option<usize>, ApiError> {
865    query
866        .get(field)
867        .map(|s| {
868            s.parse()
869                .map_err(|_| ApiError::BadRequest(format!("'{field}' must be a whole number")))
870        })
871        .transpose()
872}
873
874fn query_uuid(query: &HashMap<String, String>, field: &str) -> Result<Option<Uuid>, ApiError> {
875    query
876        .get(field)
877        .filter(|s| !s.is_empty())
878        .map(|s| {
879            Uuid::parse_str(s)
880                .map_err(|_| ApiError::BadRequest(format!("invalid UUID for '{field}'")))
881        })
882        .transpose()
883}
884
885fn body_usize(req: &Value, field: &str) -> Result<Option<usize>, ApiError> {
886    req.get(field)
887        .map(|v| {
888            v.as_u64()
889                .map(|n| n as usize)
890                .ok_or_else(|| ApiError::BadRequest(format!("'{field}' must be a whole number")))
891        })
892        .transpose()
893}
894
895fn body_str(req: &Value, field: &str) -> Option<String> {
896    req.get(field)
897        .and_then(|v| v.as_str())
898        .filter(|s| !s.is_empty())
899        .map(String::from)
900}
901
902fn string_array(req: &Value, field: &str) -> Result<Vec<String>, ApiError> {
903    match req.get(field) {
904        None | Some(Value::Null) => Ok(Vec::new()),
905        Some(Value::Array(items)) => items
906            .iter()
907            .map(|v| {
908                v.as_str()
909                    .map(String::from)
910                    .ok_or_else(|| ApiError::BadRequest(format!("'{field}' must hold strings")))
911            })
912            .collect(),
913        Some(_) => Err(ApiError::BadRequest(format!("'{field}' must be an array"))),
914    }
915}
916
917/// The query vector: the raw `embedding` when sent, else the `query` text
918/// embedded by the server.
919async fn query_embedding(db: &Db, req: &Value, text: Option<&str>) -> Result<Vec<f32>, ApiError> {
920    match (req.get("embedding"), text) {
921        (Some(_), _) => parse_embedding(req, true),
922        (None, Some(text)) => embed_query(db, text).await,
923        (None, None) => Err(ApiError::BadRequest(
924            "send 'query' text or an 'embedding' vector".into(),
925        )),
926    }
927}
928
929/// The agent the token is bound to, unless it is an admin token.
930fn token_agent(agent: &Caller) -> Result<Option<AgentId>, ApiError> {
931    match agent {
932        Some(Extension(authed)) if !authed.admin => authed.agent(),
933        _ => Ok(None),
934    }
935}
936
937/// The agent a read is scoped to: the requested one (which must match a
938/// token bound to an agent), else the token's agent, else none.
939fn scope_agent(agent: &Caller, requested: Option<Uuid>) -> Result<Option<AgentId>, ApiError> {
940    match requested {
941        Some(id) => {
942            check_body_agent(agent, AgentId(id))?;
943            Ok(Some(AgentId(id)))
944        }
945        None => token_agent(agent),
946    }
947}
948
949/// Whether a caller bound to `owner` may see `node`: its own memories and
950/// shared ones owned by no agent.
951fn visible(owner: Option<AgentId>, node: &MemoryNode) -> bool {
952    owner.is_none_or(|a| node.agent_id == a || node.agent_id.is_nil())
953}
954
955/// Load a memory and check the caller may use it.
956fn load_owned(agent: &Caller, db: &Db, id: MemoryId) -> Result<MemoryNode, ApiError> {
957    let node = db
958        .get_memory(id)
959        .map_err(|_| ApiError::NotFound(format!("memory {id} not found")))?;
960    check_owner(agent, &node)?;
961    Ok(node)
962}
963
964/// Load `ids`, keep what the caller may see, newest first, at most `limit`.
965fn newest_visible(
966    db: &mentedb::MenteDb,
967    ids: Vec<MemoryId>,
968    owner: Option<AgentId>,
969    limit: usize,
970) -> Vec<Value> {
971    let mut nodes: Vec<MemoryNode> = ids
972        .into_iter()
973        .filter_map(|id| db.get_memory(id).ok())
974        .filter(|n| visible(owner, n))
975        .collect();
976    nodes.sort_by_key(|n| std::cmp::Reverse(n.created_at));
977    nodes.iter().take(limit).map(memory_node_to_json).collect()
978}
979
980/// A memory as JSON with extra fields merged in.
981fn with_fields(node: &MemoryNode, extra: Value) -> Value {
982    let mut out = memory_node_to_json(node);
983    if let (Value::Object(o), Value::Object(e)) = (&mut out, extra) {
984        o.extend(e);
985    }
986    out
987}
988
989fn stored_edge(source: MemoryId, target: MemoryId, e: &mentedb::graph::StoredEdge) -> MemoryEdge {
990    MemoryEdge {
991        source,
992        target,
993        edge_type: e.edge_type,
994        weight: e.weight,
995        created_at: e.created_at,
996        valid_from: e.valid_from,
997        valid_until: e.valid_until,
998        label: e.label.clone(),
999    }
1000}
1001
1002fn edge_type_name(t: EdgeType) -> String {
1003    format!("{t:?}").to_lowercase()
1004}
1005
1006fn edge_to_json(e: &MemoryEdge) -> Value {
1007    json!({
1008        "source": e.source.to_string(),
1009        "target": e.target.to_string(),
1010        "edge_type": edge_type_name(e.edge_type),
1011        "weight": e.weight,
1012        "created_at": e.created_at,
1013        "valid_from": e.valid_from,
1014        "valid_until": e.valid_until,
1015        "label": e.label,
1016    })
1017}