Skip to content

Commit f37fea2

Browse files
committed
feat(isolation): add user_id and session_id support for multi-tenant memory isolation
- Add user_id and session_id fields to MemoryEntry model with optional serialization - Extend MemeBuilder with user_id() and session_id() configuration methods - Add Scope struct to VectorStore for filtering queries by user/session context - Update all CLI commands (add, ask, export, import, list) with --user-id and --session-id flags - Implement automatic scope injection in Meme::store_entries and Meme::import
1 parent 0ce557f commit f37fea2

13 files changed

Lines changed: 457 additions & 74 deletions

File tree

.github/workflows/rust-publish.yml

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,4 @@ jobs:
3636
env:
3737
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
3838
run: |
39-
for pkg in meme meme-cli; do
40-
echo "Publishing $pkg..."
41-
cargo publish -p "$pkg"
42-
echo "Waiting for crates.io to index $pkg..."
43-
sleep 30
44-
done
39+
cargo publish

meme-cli/src/commands/add.rs

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,14 @@ pub struct AddCmd {
1919
/// Import dialogues from a JSONL file.
2020
#[arg(long, value_name = "FILE")]
2121
pub file: Option<String>,
22+
23+
/// User identifier for memory isolation.
24+
#[arg(long)]
25+
pub user_id: Option<String>,
26+
27+
/// Session identifier for memory isolation.
28+
#[arg(long)]
29+
pub session_id: Option<String>,
2230
}
2331

2432
impl AddCmd {
@@ -54,10 +62,14 @@ impl AddCmd {
5462
})
5563
.transpose()?;
5664

57-
let meme = meme::MemeBuilder::new()
58-
.build()
59-
.await
60-
.map_err(|e| anyhow::anyhow!("{e}"))?;
65+
let mut builder = meme::MemeBuilder::new();
66+
if let Some(uid) = &self.user_id {
67+
builder = builder.user_id(uid);
68+
}
69+
if let Some(sid) = &self.session_id {
70+
builder = builder.session_id(sid);
71+
}
72+
let meme = builder.build().await.map_err(|e| anyhow::anyhow!("{e}"))?;
6173
meme.add_dialogue(speaker, content, timestamp)
6274
.await
6375
.map_err(|e| anyhow::anyhow!("{e}"))?;
@@ -96,10 +108,14 @@ impl AddCmd {
96108
}
97109

98110
let count = dialogues.len();
99-
let meme = meme::MemeBuilder::new()
100-
.build()
101-
.await
102-
.map_err(|e| anyhow::anyhow!("{e}"))?;
111+
let mut builder = meme::MemeBuilder::new();
112+
if let Some(uid) = &self.user_id {
113+
builder = builder.user_id(uid);
114+
}
115+
if let Some(sid) = &self.session_id {
116+
builder = builder.session_id(sid);
117+
}
118+
let meme = builder.build().await.map_err(|e| anyhow::anyhow!("{e}"))?;
103119
meme.add_dialogues(dialogues)
104120
.await
105121
.map_err(|e| anyhow::anyhow!("{e}"))?;

meme-cli/src/commands/ask.rs

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,14 @@ use clap::Args;
77
pub struct AskCmd {
88
/// The question to ask.
99
pub question: String,
10+
11+
/// User identifier for memory isolation.
12+
#[arg(long)]
13+
pub user_id: Option<String>,
14+
15+
/// Session identifier for memory isolation.
16+
#[arg(long)]
17+
pub session_id: Option<String>,
1018
}
1119

1220
impl AskCmd {
@@ -16,10 +24,14 @@ impl AskCmd {
1624
///
1725
/// Returns an error if the query fails.
1826
pub async fn run(&self) -> anyhow::Result<()> {
19-
let meme = meme::MemeBuilder::new()
20-
.build()
21-
.await
22-
.map_err(|e| anyhow::anyhow!("{e}"))?;
27+
let mut builder = meme::MemeBuilder::new();
28+
if let Some(uid) = &self.user_id {
29+
builder = builder.user_id(uid);
30+
}
31+
if let Some(sid) = &self.session_id {
32+
builder = builder.session_id(sid);
33+
}
34+
let meme = builder.build().await.map_err(|e| anyhow::anyhow!("{e}"))?;
2335

2436
let answer = meme
2537
.ask(&self.question)

meme-cli/src/commands/export.rs

Lines changed: 42 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,14 @@ pub struct ExportCmd {
88
/// Output file path (stdout if not specified).
99
#[arg(short, long)]
1010
pub output: Option<String>,
11+
12+
/// User identifier for memory isolation.
13+
#[arg(long)]
14+
pub user_id: Option<String>,
15+
16+
/// Session identifier for memory isolation.
17+
#[arg(long)]
18+
pub session_id: Option<String>,
1119
}
1220

1321
impl ExportCmd {
@@ -17,10 +25,14 @@ impl ExportCmd {
1725
///
1826
/// Returns an error if the export fails.
1927
pub async fn run(&self) -> anyhow::Result<()> {
20-
let meme = meme::MemeBuilder::new()
21-
.build()
22-
.await
23-
.map_err(|e| anyhow::anyhow!("{e}"))?;
28+
let mut builder = meme::MemeBuilder::new();
29+
if let Some(uid) = &self.user_id {
30+
builder = builder.user_id(uid);
31+
}
32+
if let Some(sid) = &self.session_id {
33+
builder = builder.session_id(sid);
34+
}
35+
let meme = builder.build().await.map_err(|e| anyhow::anyhow!("{e}"))?;
2436

2537
let entries = meme
2638
.get_all_memories()
@@ -45,6 +57,14 @@ impl ExportCmd {
4557
pub struct ImportCmd {
4658
/// Input file path.
4759
pub file: String,
60+
61+
/// User identifier for memory isolation.
62+
#[arg(long)]
63+
pub user_id: Option<String>,
64+
65+
/// Session identifier for memory isolation.
66+
#[arg(long)]
67+
pub session_id: Option<String>,
4868
}
4969

5070
impl ImportCmd {
@@ -53,16 +73,29 @@ impl ImportCmd {
5373
/// # Errors
5474
///
5575
/// Returns an error if the import fails.
56-
#[allow(clippy::unused_async)]
5776
pub async fn run(&self) -> anyhow::Result<()> {
5877
let content = std::fs::read_to_string(&self.file)?;
59-
let entries: Vec<meme::model::MemoryEntry> = serde_json::from_str(&content)?;
78+
let mut entries: Vec<meme::model::MemoryEntry> = serde_json::from_str(&content)?;
6079

6180
let count = entries.len();
62-
println!("Parsed {count} entries from {}", self.file);
63-
println!("Note: Direct entry import requires embedding recomputation.");
64-
println!("Use `meme add --file` for JSONL dialogue import instead.");
81+
println!(
82+
"Importing {count} entries from {} (recomputing embeddings)...",
83+
self.file
84+
);
85+
86+
let mut builder = meme::MemeBuilder::new();
87+
if let Some(uid) = &self.user_id {
88+
builder = builder.user_id(uid);
89+
}
90+
if let Some(sid) = &self.session_id {
91+
builder = builder.session_id(sid);
92+
}
93+
let meme = builder.build().await.map_err(|e| anyhow::anyhow!("{e}"))?;
94+
meme.import_entries(&mut entries)
95+
.await
96+
.map_err(|e| anyhow::anyhow!("{e}"))?;
6597

98+
println!("Imported {count} entries successfully.");
6699
Ok(())
67100
}
68101
}

meme-cli/src/commands/list.rs

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,14 @@ pub struct ListCmd {
1313
/// Output as JSON.
1414
#[arg(long)]
1515
pub json: bool,
16+
17+
/// User identifier for memory isolation.
18+
#[arg(long)]
19+
pub user_id: Option<String>,
20+
21+
/// Session identifier for memory isolation.
22+
#[arg(long)]
23+
pub session_id: Option<String>,
1624
}
1725

1826
impl ListCmd {
@@ -22,10 +30,14 @@ impl ListCmd {
2230
///
2331
/// Returns an error if the query fails.
2432
pub async fn run(&self) -> anyhow::Result<()> {
25-
let meme = meme::MemeBuilder::new()
26-
.build()
27-
.await
28-
.map_err(|e| anyhow::anyhow!("{e}"))?;
33+
let mut builder = meme::MemeBuilder::new();
34+
if let Some(uid) = &self.user_id {
35+
builder = builder.user_id(uid);
36+
}
37+
if let Some(sid) = &self.session_id {
38+
builder = builder.session_id(sid);
39+
}
40+
let meme = builder.build().await.map_err(|e| anyhow::anyhow!("{e}"))?;
2941

3042
let entries = meme
3143
.get_all_memories()

meme/src/embedding/api.rs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ pub struct ApiEmbedding {
1212
api_key: String,
1313
model: String,
1414
dimension: usize,
15+
max_retries: u32,
1516
}
1617

1718
impl ApiEmbedding {
@@ -29,6 +30,7 @@ impl ApiEmbedding {
2930
api_key: api_key.into(),
3031
model: model.into(),
3132
dimension,
33+
max_retries: 3,
3234
}
3335
}
3436

@@ -86,6 +88,25 @@ impl ApiEmbedding {
8688
}
8789

8890
async fn embed(&self, input: Vec<String>) -> Result<Vec<Vec<f32>>> {
91+
let mut last_err = None;
92+
for attempt in 0..self.max_retries {
93+
match self.call_embed_api(&input).await {
94+
Ok(vectors) => return Ok(vectors),
95+
Err(e) => {
96+
tracing::warn!(attempt = attempt + 1, error = %e, "embedding API call failed");
97+
last_err = Some(e);
98+
if attempt + 1 < self.max_retries {
99+
let wait = 1u64 << attempt;
100+
tokio::time::sleep(std::time::Duration::from_secs(wait)).await;
101+
}
102+
}
103+
}
104+
}
105+
Err(last_err
106+
.unwrap_or_else(|| Error::Embedding("all embedding retries exhausted".to_owned())))
107+
}
108+
109+
async fn call_embed_api(&self, input: &[String]) -> Result<Vec<Vec<f32>>> {
89110
let url = format!("{}/embeddings", self.base_url);
90111

91112
let body = serde_json::json!({

0 commit comments

Comments
 (0)