mem_types/
traits.rs

1//! Traits for MemCube and storage backends.
2
3use crate::{
4    ApiAddRequest, ApiSearchRequest, AuditEvent, AuditListOptions, ForgetMemoryRequest,
5    ForgetMemoryResponse, GetMemoryRequest, GetMemoryResponse, GraphDirection, GraphNeighbor,
6    GraphNeighborsRequest, GraphNeighborsResponse, GraphPath, GraphPathRequest, GraphPathResponse,
7    GraphPathsRequest, GraphPathsResponse, MemoryEdge, MemoryNode, MemoryResponse, SearchResponse,
8    UpdateMemoryRequest, UpdateMemoryResponse,
9};
10use async_trait::async_trait;
11use std::collections::HashMap;
12
13/// Result of a vector search hit (id + score).
14#[derive(Debug, Clone)]
15pub struct VecSearchHit {
16    pub id: String,
17    pub score: f64,
18}
19
20/// Graph store abstraction (subset of MemOS BaseGraphDB).
21#[async_trait]
22pub trait GraphStore: Send + Sync {
23    /// Add a single memory node.
24    async fn add_node(
25        &self,
26        id: &str,
27        memory: &str,
28        metadata: &HashMap<String, serde_json::Value>,
29        user_name: Option<&str>,
30    ) -> Result<(), GraphStoreError>;
31
32    /// Add multiple nodes in batch.
33    async fn add_nodes_batch(
34        &self,
35        nodes: &[MemoryNode],
36        user_name: Option<&str>,
37    ) -> Result<(), GraphStoreError>;
38
39    /// Add multiple edges in batch.
40    async fn add_edges_batch(
41        &self,
42        edges: &[MemoryEdge],
43        user_name: Option<&str>,
44    ) -> Result<(), GraphStoreError>;
45
46    /// Get one node by id.
47    async fn get_node(
48        &self,
49        id: &str,
50        include_embedding: bool,
51    ) -> Result<Option<MemoryNode>, GraphStoreError>;
52
53    /// Get multiple nodes by ids.
54    async fn get_nodes(
55        &self,
56        ids: &[String],
57        include_embedding: bool,
58    ) -> Result<Vec<MemoryNode>, GraphStoreError>;
59
60    /// Get neighbors of one node, optionally filtered by relation and direction.
61    async fn get_neighbors(
62        &self,
63        id: &str,
64        relation: Option<&str>,
65        direction: GraphDirection,
66        limit: usize,
67        include_embedding: bool,
68        user_name: Option<&str>,
69    ) -> Result<Vec<GraphNeighbor>, GraphStoreError>;
70
71    /// Shortest path query between source and target by BFS hops.
72    async fn shortest_path(
73        &self,
74        source_id: &str,
75        target_id: &str,
76        relation: Option<&str>,
77        direction: GraphDirection,
78        max_depth: usize,
79        include_deleted: bool,
80        user_name: Option<&str>,
81    ) -> Result<Option<GraphPath>, GraphStoreError>;
82
83    /// Enumerate top-k shortest simple paths by BFS hops.
84    async fn find_paths(
85        &self,
86        source_id: &str,
87        target_id: &str,
88        relation: Option<&str>,
89        direction: GraphDirection,
90        max_depth: usize,
91        top_k: usize,
92        include_deleted: bool,
93        user_name: Option<&str>,
94    ) -> Result<Vec<GraphPath>, GraphStoreError>;
95
96    /// Search by embedding vector (returns node ids + scores).
97    async fn search_by_embedding(
98        &self,
99        vector: &[f32],
100        top_k: usize,
101        user_name: Option<&str>,
102    ) -> Result<Vec<VecSearchHit>, GraphStoreError>;
103
104    /// Get all memory items for a scope and user.
105    async fn get_all_memory_items(
106        &self,
107        scope: &str,
108        user_name: &str,
109        include_embedding: bool,
110    ) -> Result<Vec<MemoryNode>, GraphStoreError>;
111
112    /// Update fields of an existing node (memory and/or metadata).
113    async fn update_node(
114        &self,
115        id: &str,
116        fields: &HashMap<String, serde_json::Value>,
117        user_name: Option<&str>,
118    ) -> Result<(), GraphStoreError>;
119
120    /// Delete a node (hard delete). If `user_name` is `Some`, implementation must verify
121    /// the node belongs to that user/cube (e.g. via metadata) before deleting; return error if not owner.
122    async fn delete_node(&self, id: &str, user_name: Option<&str>) -> Result<(), GraphStoreError>;
123
124    /// Delete all edges connected to a node. Returns number of deleted edges.
125    async fn delete_edges_by_node(
126        &self,
127        id: &str,
128        user_name: Option<&str>,
129    ) -> Result<usize, GraphStoreError>;
130}
131
132/// Vector store abstraction (subset of MemOS BaseVecDB).
133#[async_trait]
134pub trait VecStore: Send + Sync {
135    /// Add items (id, vector, payload).
136    async fn add(
137        &self,
138        items: &[VecStoreItem],
139        collection: Option<&str>,
140    ) -> Result<(), VecStoreError>;
141
142    /// Search by vector.
143    async fn search(
144        &self,
145        query_vector: &[f32],
146        top_k: usize,
147        filter: Option<&HashMap<String, serde_json::Value>>,
148        collection: Option<&str>,
149    ) -> Result<Vec<VecSearchHit>, VecStoreError>;
150
151    /// Get by ids.
152    async fn get_by_ids(
153        &self,
154        ids: &[String],
155        collection: Option<&str>,
156    ) -> Result<Vec<VecStoreItem>, VecStoreError>;
157
158    /// Delete by ids.
159    async fn delete(&self, ids: &[String], collection: Option<&str>) -> Result<(), VecStoreError>;
160
161    /// Upsert items: insert or replace by id. Ensures full payload and avoids delete+add window.
162    async fn upsert(
163        &self,
164        items: &[VecStoreItem],
165        collection: Option<&str>,
166    ) -> Result<(), VecStoreError>;
167}
168
169/// Item for vector store (id, vector, payload).
170#[derive(Clone, Debug)]
171pub struct VecStoreItem {
172    pub id: String,
173    pub vector: Vec<f32>,
174    pub payload: HashMap<String, serde_json::Value>,
175}
176
177/// Embedder: text -> vector(s).
178#[async_trait]
179pub trait Embedder: Send + Sync {
180    /// Embed a single text. Default implementation uses embed_batch.
181    async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedderError> {
182        let v = self.embed_batch(&[text.to_string()]).await?;
183        v.into_iter().next().ok_or(EmbedderError::EmptyResponse)
184    }
185
186    /// Embed multiple texts.
187    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbedderError>;
188}
189
190/// MemCube abstraction: add, search, update, and forget memories.
191#[async_trait]
192pub trait MemCube: Send + Sync {
193    /// Add memories from request; returns MemoryResponse.
194    async fn add_memories(&self, req: &ApiAddRequest) -> Result<MemoryResponse, MemCubeError>;
195
196    /// Search memories from request; returns SearchResponse.
197    async fn search_memories(&self, req: &ApiSearchRequest)
198        -> Result<SearchResponse, MemCubeError>;
199
200    /// Update an existing memory (partial fields); re-embeds if memory text changed.
201    async fn update_memory(
202        &self,
203        req: &UpdateMemoryRequest,
204    ) -> Result<UpdateMemoryResponse, MemCubeError>;
205
206    /// Forget (soft or hard delete) a memory.
207    async fn forget_memory(
208        &self,
209        req: &ForgetMemoryRequest,
210    ) -> Result<ForgetMemoryResponse, MemCubeError>;
211
212    /// Get a single memory by id (within user/cube scope).
213    async fn get_memory(&self, req: &GetMemoryRequest) -> Result<GetMemoryResponse, MemCubeError>;
214
215    /// Query graph neighbors for one memory id.
216    async fn graph_neighbors(
217        &self,
218        req: &GraphNeighborsRequest,
219    ) -> Result<GraphNeighborsResponse, MemCubeError>;
220
221    /// Query shortest path between two memory nodes.
222    async fn graph_path(&self, req: &GraphPathRequest) -> Result<GraphPathResponse, MemCubeError>;
223
224    /// Query top-k shortest paths between two memory nodes.
225    async fn graph_paths(
226        &self,
227        req: &GraphPathsRequest,
228    ) -> Result<GraphPathsResponse, MemCubeError>;
229}
230
231#[derive(Debug, thiserror::Error)]
232pub enum GraphStoreError {
233    #[error("graph store error: {0}")]
234    Other(String),
235}
236
237#[derive(Debug, thiserror::Error)]
238pub enum VecStoreError {
239    #[error("vector store error: {0}")]
240    Other(String),
241}
242
243#[derive(Debug, thiserror::Error)]
244pub enum EmbedderError {
245    #[error("embedder error: {0}")]
246    Other(String),
247    #[error("empty response")]
248    EmptyResponse,
249}
250
251/// Audit event store: append-only log with optional list by user/cube/time and pagination.
252#[async_trait]
253pub trait AuditStore: Send + Sync {
254    /// Append one audit event.
255    async fn append(&self, event: AuditEvent) -> Result<(), AuditStoreError>;
256
257    /// List events with optional filters and limit/offset. Newest first.
258    async fn list(&self, opts: &AuditListOptions) -> Result<Vec<AuditEvent>, AuditStoreError>;
259}
260
261#[derive(Debug, thiserror::Error)]
262pub enum AuditStoreError {
263    #[error("audit store error: {0}")]
264    Other(String),
265}
266
267#[derive(Debug, thiserror::Error)]
268pub enum MemCubeError {
269    #[error("mem cube error: {0}")]
270    Other(String),
271    #[error("bad request: {0}")]
272    BadRequest(String),
273    #[error("not found: {0}")]
274    NotFound(String),
275    #[error("embedder: {0}")]
276    Embedder(#[from] EmbedderError),
277    #[error("graph: {0}")]
278    Graph(#[from] GraphStoreError),
279    #[error("vector: {0}")]
280    Vec(#[from] VecStoreError),
281}
mem_types/traits.rs

mem_types/
traits.rs
