use std::collections::{HashMap, HashSet};
use rusqlite::{Connection, params, params_from_iter, types::Value};
use crate::engine::memory::entity::{EntityCandidate, normalize_entity_value};
use crate::engine::memory::types::{
ClaimRelationship, EntityReference, MatchKind, MatchReason, MemoryError, MemoryStatus,
MemoryType, RecallHit, RelationshipKind,
};
use crate::engine::model::{EMBEDDING_MODEL_ID, Embedding};
use super::util::RELATED_CANDIDATE_LIMIT;
pub(super) fn entity_related_sql(seed_count: usize) -> String {
let placeholders = (1..=seed_count)
.map(|index| format!("?{index}"))
.collect::<Vec<_>>()
.join(", ");
let project_idx = seed_count + 1;
let type_idx = seed_count + 2;
let history_idx = seed_count + 3;
let limit_idx = seed_count + 4;
format!(
"SELECT claims.id, claims.memory_type, claims.statement,
claims.keywords, projects.path, claims.valid_from,
claims.valid_until, claims.created_at, claims.status,
claims.superseded_by,
count(DISTINCT bridge.entity_id) AS shared,
group_concat(entities.kind || char(31) || entities.value, char(30)) AS shared_entities
FROM claim_entities seed
JOIN claim_entities bridge
ON bridge.entity_id = seed.entity_id
AND bridge.claim_id != seed.claim_id
JOIN claims ON claims.id = bridge.claim_id
JOIN projects ON projects.id = claims.project_id
JOIN entities ON entities.id = bridge.entity_id
WHERE seed.claim_id IN ({placeholders})
AND (?{project_idx} IS NULL OR projects.path = ?{project_idx})
AND (?{type_idx} IS NULL OR claims.memory_type = ?{type_idx})
AND (?{history_idx} OR claims.status = 'active')
GROUP BY claims.id
ORDER BY shared DESC, claims.created_at DESC
LIMIT ?{limit_idx}"
)
}
pub(super) struct RecallFill<'a> {
pub(super) hits: &'a mut Vec<RecallHit>,
pub(super) used_tokens: &'a mut usize,
pub(super) limit: usize,
pub(super) max_tokens: usize,
}
pub(super) struct RelationCandidate {
pub(super) claim_id: String,
pub(super) seed_display_id: String,
pub(super) seed_score: f64,
pub(super) kind: RelationshipKind,
pub(super) rationale: String,
}
pub(super) fn relation_candidates(
hits: &[RecallHit],
history: bool,
memory_type: Option<MemoryType>,
) -> Vec<RelationCandidate> {
let mut candidates = hits
.iter()
.filter(|hit| {
hit.match_reasons.iter().any(|reason| {
matches!(
reason.kind,
MatchKind::Lexical | MatchKind::Entity | MatchKind::Vector
)
})
})
.flat_map(|hit| {
hit.relationships.iter().filter_map(|relationship| {
if memory_type.is_some_and(|kind| relationship.memory_type != kind)
|| relationship.kind == RelationshipKind::Duplicates
|| (relationship.kind == RelationshipKind::Revises && !history)
{
return None;
}
Some(RelationCandidate {
claim_id: relationship.claim_id.strip_prefix("mem_")?.to_string(),
seed_display_id: hit.display_id.clone(),
seed_score: hit.score,
kind: relationship.kind,
rationale: relationship.rationale.clone(),
})
})
})
.collect::<Vec<_>>();
candidates.sort_by(|left, right| {
relation_priority(left.kind)
.cmp(&relation_priority(right.kind))
.then_with(|| right.seed_score.total_cmp(&left.seed_score))
.then_with(|| left.claim_id.cmp(&right.claim_id))
});
let mut candidate_ids = HashSet::new();
candidates.retain(|candidate| candidate_ids.insert(candidate.claim_id.clone()));
candidates.truncate(RELATED_CANDIDATE_LIMIT);
candidates
}
pub(super) fn recall_match_reasons(
kind: MatchKind,
detail: impl Into<String>,
status: MemoryStatus,
relationships: &[ClaimRelationship],
) -> Vec<MatchReason> {
let mut reasons = vec![MatchReason {
kind,
detail: detail.into(),
}];
if status == MemoryStatus::Superseded {
reasons.push(MatchReason {
kind: MatchKind::History,
detail: "superseded claim included by history search".to_string(),
});
}
if status == MemoryStatus::Active {
for relationship in relationships.iter().filter(|relationship| {
relationship.kind == RelationshipKind::Contradicts
&& relationship.status == MemoryStatus::Active
}) {
let detail = relation_match_detail(
relationship.kind,
&relationship.display_id,
&relationship.rationale,
);
if !reasons
.iter()
.any(|reason| reason.kind == MatchKind::Relation && reason.detail == detail)
{
reasons.push(MatchReason {
kind: MatchKind::Relation,
detail,
});
}
}
}
reasons
}
pub(super) fn relation_match_detail(kind: RelationshipKind, peer: &str, rationale: &str) -> String {
let relation = match kind {
RelationshipKind::Contradicts => format!("OPEN CONTRADICTION: contradicts {peer}"),
RelationshipKind::Supports => format!("supports relation with {peer}"),
RelationshipKind::Revises => format!("revises relation with {peer}"),
RelationshipKind::Duplicates => format!("duplicates relation with {peer}"),
};
if rationale.trim().is_empty() {
relation
} else {
format!("{relation}: {}", rationale.trim())
}
}
pub(super) const fn relation_priority(kind: RelationshipKind) -> u8 {
match kind {
RelationshipKind::Contradicts => 0,
RelationshipKind::Revises => 1,
RelationshipKind::Supports => 2,
RelationshipKind::Duplicates => 3,
}
}
pub(super) const fn relation_score_factor(kind: RelationshipKind) -> f64 {
match kind {
RelationshipKind::Contradicts => 0.9,
RelationshipKind::Revises => 0.7,
RelationshipKind::Supports => 0.6,
RelationshipKind::Duplicates => 0.4,
}
}
pub(super) fn entity_match_detail(shared: i64, entities: &[EntityReference]) -> String {
let names = entities
.iter()
.map(|entity| format!("{}:{}", entity.kind, entity.value))
.collect::<Vec<_>>()
.join(", ");
if names.is_empty() {
format!("{shared} shared entities")
} else {
format!("{shared} shared entities: {names}")
}
}
pub(super) struct EntityRelatedRow {
pub(super) id: String,
pub(super) memory_type: MemoryType,
pub(super) statement: String,
pub(super) keywords: String,
pub(super) project: String,
pub(super) valid_from: Option<i64>,
pub(super) valid_until: Option<i64>,
pub(super) status: MemoryStatus,
pub(super) superseded_by: Option<String>,
pub(super) shared: i64,
pub(super) related_by: Vec<EntityReference>,
}
pub(super) fn load_entity_related_rows(
conn: &Connection,
seed_ids: &[String],
project: Option<String>,
memory_type: Option<&str>,
history: bool,
limit: usize,
) -> anyhow::Result<Vec<EntityRelatedRow>> {
if seed_ids.is_empty() {
return Ok(Vec::new());
}
let sql = entity_related_sql(seed_ids.len());
let mut stmt = conn.prepare(&sql)?;
let mut params: Vec<rusqlite::types::Value> = seed_ids
.iter()
.map(|id| rusqlite::types::Value::Text(id.clone()))
.collect();
params.push(match project {
Some(path) => rusqlite::types::Value::Text(path),
None => rusqlite::types::Value::Null,
});
params.push(match memory_type {
Some(value) => rusqlite::types::Value::Text(value.to_string()),
None => rusqlite::types::Value::Null,
});
params.push(rusqlite::types::Value::Integer(i64::from(history)));
params.push(rusqlite::types::Value::Integer(i64::try_from(
RELATED_CANDIDATE_LIMIT.min(limit.saturating_mul(2)),
)?));
let rows = stmt
.query_map(rusqlite::params_from_iter(params), |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, String>(4)?,
row.get::<_, Option<i64>>(5)?,
row.get::<_, Option<i64>>(6)?,
row.get::<_, i64>(7)?,
row.get::<_, String>(8)?,
row.get::<_, Option<String>>(9)?,
row.get::<_, i64>(10)?,
row.get::<_, Option<String>>(11)?,
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
let mut out = Vec::new();
for (
id,
raw_type,
statement,
keywords,
project_path,
valid_from,
valid_until,
_created_at,
raw_status,
superseded_by,
shared,
shared_entities,
) in rows
{
let memory_type: MemoryType = raw_type.parse()?;
let status: MemoryStatus = raw_status.parse()?;
let related_by = parse_shared_entities(shared_entities.as_deref().unwrap_or(""));
if related_by.is_empty() {
continue;
}
out.push(EntityRelatedRow {
id,
memory_type,
statement,
keywords,
project: project_path,
valid_from,
valid_until,
status,
superseded_by,
shared,
related_by,
});
}
Ok(out)
}
pub(super) fn query_entity_sql(pair_count: usize) -> String {
let pair_clauses = (0..pair_count)
.map(|index| {
let kind_idx = index * 2 + 1;
let norm_idx = index * 2 + 2;
format!("(entities.kind = ?{kind_idx} AND entities.normalized = ?{norm_idx})")
})
.collect::<Vec<_>>()
.join(" OR ");
let project_idx = pair_count * 2 + 1;
let type_idx = pair_count * 2 + 2;
let history_idx = pair_count * 2 + 3;
let limit_idx = pair_count * 2 + 4;
format!(
"SELECT claims.id, claims.memory_type, claims.statement,
claims.keywords, projects.path, claims.valid_from,
claims.valid_until, claims.created_at, claims.status,
claims.superseded_by,
count(DISTINCT entities.id) AS shared,
group_concat(entities.kind || char(31) || entities.value, char(30)) AS shared_entities
FROM entities
JOIN claim_entities ON claim_entities.entity_id = entities.id
JOIN claims ON claims.id = claim_entities.claim_id
JOIN projects ON projects.id = claims.project_id
WHERE ({pair_clauses})
AND (?{project_idx} IS NULL OR projects.path = ?{project_idx})
AND (?{type_idx} IS NULL OR claims.memory_type = ?{type_idx})
AND (?{history_idx} OR claims.status = 'active')
GROUP BY claims.id
ORDER BY shared DESC, claims.created_at DESC
LIMIT ?{limit_idx}"
)
}
pub(super) fn load_query_entity_rows(
conn: &Connection,
query_entities: &[EntityCandidate],
project: Option<String>,
memory_type: Option<&str>,
history: bool,
limit: usize,
) -> anyhow::Result<Vec<EntityRelatedRow>> {
if query_entities.is_empty() {
return Ok(Vec::new());
}
let mut pairs: Vec<(&str, String)> = Vec::new();
let mut seen = HashSet::new();
for entity in query_entities {
let normalized = normalize_entity_value(&entity.value);
let key = format!("{}\0{normalized}", entity.kind);
if !seen.insert(key) {
continue;
}
pairs.push((entity.kind, normalized));
}
if pairs.is_empty() {
return Ok(Vec::new());
}
let sql = query_entity_sql(pairs.len());
let mut stmt = conn.prepare(&sql)?;
let mut params: Vec<rusqlite::types::Value> = Vec::new();
for (kind, normalized) in &pairs {
params.push(rusqlite::types::Value::Text((*kind).to_string()));
params.push(rusqlite::types::Value::Text(normalized.clone()));
}
params.push(match project {
Some(path) => rusqlite::types::Value::Text(path),
None => rusqlite::types::Value::Null,
});
params.push(match memory_type {
Some(value) => rusqlite::types::Value::Text(value.to_string()),
None => rusqlite::types::Value::Null,
});
params.push(rusqlite::types::Value::Integer(i64::from(history)));
params.push(rusqlite::types::Value::Integer(i64::try_from(
RELATED_CANDIDATE_LIMIT.min(limit.saturating_mul(2)),
)?));
let rows = stmt
.query_map(rusqlite::params_from_iter(params), |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
row.get::<_, String>(4)?,
row.get::<_, Option<i64>>(5)?,
row.get::<_, Option<i64>>(6)?,
row.get::<_, i64>(7)?,
row.get::<_, String>(8)?,
row.get::<_, Option<String>>(9)?,
row.get::<_, i64>(10)?,
row.get::<_, Option<String>>(11)?,
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
let mut out = Vec::new();
for (
id,
raw_type,
statement,
keywords,
project_path,
valid_from,
valid_until,
_created_at,
raw_status,
superseded_by,
shared,
shared_entities,
) in rows
{
let memory_type: MemoryType = raw_type.parse()?;
let status: MemoryStatus = raw_status.parse()?;
let related_by = parse_query_entities(shared_entities.as_deref().unwrap_or(""));
if related_by.is_empty() {
continue;
}
out.push(EntityRelatedRow {
id,
memory_type,
statement,
keywords,
project: project_path,
valid_from,
valid_until,
status,
superseded_by,
shared,
related_by,
});
}
Ok(out)
}
pub(super) fn parse_query_entities(raw: &str) -> Vec<EntityReference> {
let mut entities = Vec::new();
let mut seen = HashSet::new();
for item in raw.split('\u{1e}').filter(|item| !item.is_empty()) {
let Some((kind, value)) = item.split_once('\u{1f}') else {
continue;
};
let key = format!("{kind}\0{value}");
if !seen.insert(key) {
continue;
}
entities.push(EntityReference {
kind: kind.to_string(),
value: value.to_string(),
origins: vec!["query".to_string()],
});
}
entities
}
pub(super) fn parse_shared_entities(raw: &str) -> Vec<EntityReference> {
let mut entities = Vec::new();
let mut seen = HashSet::new();
for item in raw.split('\u{1e}').filter(|item| !item.is_empty()) {
let Some((kind, value)) = item.split_once('\u{1f}') else {
continue;
};
let key = format!("{kind}\0{value}");
if !seen.insert(key) {
continue;
}
entities.push(EntityReference {
kind: kind.to_string(),
value: value.to_string(),
origins: vec!["shared".to_string()],
});
}
entities
}
pub(super) fn load_memory_rows(
conn: &Connection,
claim_ids: &[String],
project: Option<&str>,
memory_type: Option<&str>,
history: bool,
) -> anyhow::Result<HashMap<String, MemoryRow>> {
const BATCH_SIZE: usize = 900;
let mut memory_rows = HashMap::new();
for batch in claim_ids.chunks(BATCH_SIZE) {
let placeholders = (1..=batch.len())
.map(|index| format!("?{index}"))
.collect::<Vec<_>>()
.join(",");
let project_index = batch.len() + 1;
let memory_type_index = batch.len() + 2;
let history_index = batch.len() + 3;
let sql = format!(
r"SELECT claims.id, claims.memory_type, claims.statement,
claims.keywords, projects.path, claims.valid_from,
claims.valid_until, claims.created_at, claims.status,
claims.superseded_by, 0.0
FROM claims
JOIN projects ON projects.id = claims.project_id
WHERE claims.id IN ({placeholders})
AND (?{project_index} IS NULL OR projects.path = ?{project_index})
AND (?{memory_type_index} IS NULL OR claims.memory_type = ?{memory_type_index})
AND (?{history_index} OR claims.status = 'active')"
);
let mut values = batch.iter().cloned().map(Value::Text).collect::<Vec<_>>();
values.push(project.map_or(Value::Null, |value| Value::Text(value.to_string())));
values.push(memory_type.map_or(Value::Null, |value| Value::Text(value.to_string())));
values.push(Value::Integer(i64::from(history)));
let mut stmt = conn.prepare(&sql)?;
let rows = stmt.query_map(params_from_iter(values), map_memory_row)?;
for row in rows {
let row = row?;
memory_rows.insert(row.id.clone(), row);
}
}
Ok(memory_rows)
}
pub(super) struct MemoryRow {
pub(super) id: String,
pub(super) memory_type: MemoryType,
pub(super) statement: String,
pub(super) keywords: String,
pub(super) project: String,
pub(super) valid_from: Option<i64>,
pub(super) valid_until: Option<i64>,
pub(super) created_at: i64,
pub(super) status: MemoryStatus,
pub(super) superseded_by: Option<String>,
pub(super) rank: f64,
}
pub(super) struct SemanticCandidate {
pub(super) row: MemoryRow,
pub(super) similarity: f64,
}
pub(super) struct MemoryEmbedding {
pub(super) claim_id: String,
pub(super) embedding: Embedding,
}
pub(super) fn map_memory_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<MemoryRow> {
let raw_type = row.get::<_, String>(1)?;
let memory_type = raw_type.parse().map_err(|error| {
rusqlite::Error::FromSqlConversionFailure(1, rusqlite::types::Type::Text, Box::new(error))
})?;
let raw_status = row.get::<_, String>(8)?;
let status = raw_status.parse().map_err(|error| {
rusqlite::Error::FromSqlConversionFailure(8, rusqlite::types::Type::Text, Box::new(error))
})?;
Ok(MemoryRow {
id: row.get(0)?,
memory_type,
statement: row.get(2)?,
keywords: row.get(3)?,
project: row.get(4)?,
valid_from: row.get(5)?,
valid_until: row.get(6)?,
created_at: row.get(7)?,
status,
superseded_by: row.get(9)?,
rank: row.get(10)?,
})
}
pub(super) fn store_embeddings(
conn: &Connection,
embeddings: &[MemoryEmbedding],
) -> anyhow::Result<()> {
if embeddings.is_empty() {
return Ok(());
}
let tx = conn.unchecked_transaction()?;
for record in embeddings {
let vector = Vec::<u8>::from(&record.embedding);
tx.execute(
"DELETE FROM claim_embeddings WHERE claim_id = ?1",
params![record.claim_id],
)?;
let inserted = tx.execute(
"INSERT INTO claim_embeddings(
claim_id, project_id, embedding_model, memory_type, memory_status, embedding
)
SELECT id, project_id, ?2, memory_type, status, ?3
FROM claims
WHERE id = ?1",
params![record.claim_id, EMBEDDING_MODEL_ID, vector],
)?;
if inserted != 1 {
return Err(MemoryError::MissingClaim(record.claim_id.clone()).into());
}
}
tx.commit()?;
Ok(())
}