mirror of
https://gitea.ingwaz.work/Ingwaz/openbrain-mcp.git
synced 2026-03-31 14:49:06 +00:00
feat: implement batch_store endpoint (Issue #12)
- Add batch_store tool accepting 1-50 entries per call - Single DB transaction for atomicity - Returns individual IDs/status per entry - Add batch_store_memories() to Database layer - Add 6 test cases - Backward compatible - existing store unchanged Expected impact: 50-60% reduction in store API calls
This commit is contained in:
33
src/db.rs
33
src/db.rs
@@ -174,3 +174,36 @@ impl Database {
|
||||
Ok(row.get("count"))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Result for a single batch entry
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BatchStoreResult {
|
||||
pub id: String,
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
impl Database {
|
||||
/// Store multiple memories in a single transaction
|
||||
pub async fn batch_store_memories(
|
||||
&self,
|
||||
agent_id: &str,
|
||||
entries: Vec<(String, Value, Vec<f32>, Vec<String>)>,
|
||||
) -> Result<Vec<BatchStoreResult>> {
|
||||
let mut client = self.pool.get().await?;
|
||||
let transaction = client.transaction().await?;
|
||||
let mut results = Vec::with_capacity(entries.len());
|
||||
|
||||
for (content, metadata, embedding, keywords) in entries {
|
||||
let id = Uuid::new_v4();
|
||||
let vector = Vector::from(embedding);
|
||||
transaction.execute(
|
||||
r#"INSERT INTO memories (id, agent_id, content, embedding, keywords, metadata) VALUES ($1, $2, $3, $4, $5, $6)"#,
|
||||
&[&id, &agent_id, &content, &vector, &keywords, &metadata],
|
||||
).await?;
|
||||
results.push(BatchStoreResult { id: id.to_string(), status: "stored".to_string() });
|
||||
}
|
||||
transaction.commit().await?;
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user