Public/openbrain-mcp
1
0
Fork 0
mirror of https://gitea.ingwaz.work/Ingwaz/openbrain-mcp.git synced 2026-03-31 14:49:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add TTL expiry for transient facts

Browse Source
This commit is contained in:
Agent Zero
2026-03-24 03:20:10 +00:00
parent 1314015479
commit 5d5c042dd1
12 changed files with 241 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
src/db.rs
View File

@@ -29,6 +29,7 @@ pub struct MemoryRecord {
pub keywords: Vec<String>,
pub metadata: serde_json::Value,
pub created_at: chrono::DateTime<chrono::Utc>,
pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
}
/// Query result with similarity score
@@ -75,6 +76,7 @@ impl Database {
embedding: &[f32],
keywords: &[String],
metadata: serde_json::Value,
expires_at: Option<chrono::DateTime<chrono::Utc>>,
) -> Result<Uuid> {
let client = self.pool.get().await?;
let id = Uuid::new_v4();
@@ -83,10 +85,10 @@ impl Database {
client
.execute(
r#"
INSERT INTO memories (id, agent_id, content, embedding, keywords, metadata)
VALUES ($1, $2, $3, $4, $5, $6)
INSERT INTO memories (id, agent_id, content, embedding, keywords, metadata, expires_at)
VALUES ($1, $2, $3, $4, $5, $6, $7)
"#,
&[&id, &agent_id, &content, &vector, &keywords, &metadata],
&[&id, &agent_id, &content, &vector, &keywords, &metadata, &expires_at],
)
.await
.context("Failed to store memory")?;
@@ -123,6 +125,7 @@ impl Database {
keywords,
metadata,
created_at,
expires_at,
(1 - (embedding <=> $1))::real AS vector_score,
CASE
WHEN search_query.query_text IS NULL THEN 0::real
@@ -133,6 +136,7 @@ impl Database {
FROM memories
CROSS JOIN search_query
WHERE memories.agent_id = $3
AND (memories.expires_at IS NULL OR memories.expires_at > NOW())
),
ranked AS (
SELECT
@@ -147,6 +151,7 @@ impl Database {
keywords,
metadata,
created_at,
expires_at,
vector_score,
text_score,
CASE
@@ -184,6 +189,7 @@ impl Database {
keywords: row.get("keywords"),
metadata: row.get("metadata"),
created_at: row.get("created_at"),
expires_at: row.get("expires_at"),
},
similarity: row.get("hybrid_score"),
vector_score: row.get("vector_score"),
@@ -224,12 +230,25 @@ impl Database {
let client = self.pool.get().await?;
let row = client
.query_one(
"SELECT COUNT(*) as count FROM memories WHERE agent_id = $1",
"SELECT COUNT(*) as count FROM memories WHERE agent_id = $1 AND (expires_at IS NULL OR expires_at > NOW())",
&[&agent_id],
)
.await?;
Ok(row.get("count"))
}
/// Delete expired memories across all agents
pub async fn cleanup_expired_memories(&self) -> Result<u64> {
let client = self.pool.get().await?;
let deleted = client
.execute(
"DELETE FROM memories WHERE expires_at IS NOT NULL AND expires_at <= NOW()",
&[],
)
.await
.context("Failed to cleanup expired memories")?;
Ok(deleted)
}
}
@@ -238,6 +257,7 @@ impl Database {
pub struct BatchStoreResult {
pub id: String,
pub status: String,
pub expires_at: Option<String>,
}
impl Database {
@@ -245,20 +265,30 @@ impl Database {
pub async fn batch_store_memories(
&self,
agent_id: &str,
entries: Vec<(String, Value, Vec<f32>, Vec<String>)>,
entries: Vec<(
String,
Value,
Vec<f32>,
Vec<String>,
Option<chrono::DateTime<chrono::Utc>>,
)>,
) -> 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 {
for (content, metadata, embedding, keywords, expires_at) 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],
r#"INSERT INTO memories (id, agent_id, content, embedding, keywords, metadata, expires_at) VALUES ($1, $2, $3, $4, $5, $6, $7)"#,
&[&id, &agent_id, &content, &vector, &keywords, &metadata, &expires_at],
).await?;
results.push(BatchStoreResult { id: id.to_string(), status: "stored".to_string() });
results.push(BatchStoreResult {
id: id.to_string(),
status: "stored".to_string(),
expires_at: expires_at.map(|ts| ts.to_rfc3339()),
});
}
transaction.commit().await?;
Ok(results)