use std::collections::HashSet;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use rusqlite::{Connection, Transaction, params};
use sha2::{Digest as _, Sha256};
use crate::engine::memory::external_id;
use crate::engine::memory::types::{
ClaimRelationship, MatchReason, MemoryError, MemoryType, MemoryTypeCounts, SourceReference,
};
pub(super) const DISPLAY_ID_CHARS: usize = 12;
pub(super) const MIN_ID_PREFIX_CHARS: usize = 8;
pub(super) const MAX_RECALL_LIMIT: usize = 100;
pub(super) const RELATED_CANDIDATE_LIMIT: usize = 50;
pub(super) const CONSOLIDATION_CANDIDATE_LIMIT: usize = 12;
const CONSOLIDATION_QUERY_TERM_LIMIT: usize = 64;
pub(super) const EMBEDDING_BACKFILL_LIMIT: usize = 64;
pub(super) const SQLITE_VEC_MAX_K: usize = 4_096;
pub(super) const SEMANTIC_MIN_SIMILARITY: f64 = 0.55;
pub(super) fn memory_type_and_text_conn(
conn: &Connection,
id: &str,
) -> anyhow::Result<(MemoryType, String)> {
let (raw_type, text): (String, String) = conn.query_row(
"SELECT memory_type, statement FROM claims WHERE id = ?1",
params![id],
|row| Ok((row.get(0)?, row.get(1)?)),
)?;
Ok((raw_type.parse()?, text))
}
pub(super) fn memory_type_and_text(
tx: &Transaction<'_>,
id: &str,
) -> anyhow::Result<(MemoryType, String)> {
memory_type_and_text_conn(tx, id)
}
pub(super) fn display_id_in_tx(tx: &Transaction<'_>, id: &str) -> anyhow::Result<String> {
let start = DISPLAY_ID_CHARS.min(id.len());
for chars in start..=id.len() {
let prefix = &id[..chars];
let matches: i64 = tx.query_row(
"SELECT count(*) FROM claims WHERE id LIKE ?1",
params![format!("{prefix}%")],
|row| row.get(0),
)?;
if matches <= 1 {
return Ok(format!("mem_{prefix}"));
}
}
Ok(external_id(id))
}
pub(super) fn memory_is_tombstoned(tx: &Transaction<'_>, id: &str) -> anyhow::Result<bool> {
Ok(tx.query_row(
"SELECT EXISTS(SELECT 1 FROM tombstones WHERE kind = 'memory' AND key = ?1)",
params![id],
|row| row.get(0),
)?)
}
pub(super) fn claim_id(project: &str, memory_type: MemoryType, statement: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(b"goosedump-memory-v2\0");
hasher.update(project.as_bytes());
hasher.update(b"\0");
hasher.update(memory_type.as_str().as_bytes());
hasher.update(b"\0");
hasher.update(statement.trim().to_lowercase().as_bytes());
format!("{:x}", hasher.finalize())
}
pub(super) fn sha256_hex(value: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(value);
format!("{:x}", hasher.finalize())
}
pub(super) fn parse_id_prefix(target: &str) -> anyhow::Result<&str> {
let prefix = target.strip_prefix("mem_").unwrap_or(target);
if prefix.len() < MIN_ID_PREFIX_CHARS || prefix.len() > 64 {
return Err(MemoryError::InvalidId(format!(
"memory ID must contain between {MIN_ID_PREFIX_CHARS} and 64 hexadecimal characters"
))
.into());
}
if !prefix.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Err(MemoryError::InvalidId(
"memory ID must be hexadecimal and may start with 'mem_'".to_string(),
)
.into());
}
Ok(prefix)
}
fn format_fts_query<'a>(terms: impl Iterator<Item = &'a str>) -> String {
terms
.map(|term| format!("\"{}\"", term.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(" OR ")
}
pub(super) fn fts_query(query: &str) -> String {
format_fts_query(query.split_whitespace().filter(|term| !term.is_empty()))
}
pub(super) fn consolidation_fts_query(query: &str) -> String {
let mut seen = HashSet::new();
format_fts_query(
query
.split_whitespace()
.filter(|term| seen.insert(term.to_lowercase()))
.take(CONSOLIDATION_QUERY_TERM_LIMIT),
)
}
pub(super) fn split_keywords(keywords: &str) -> Vec<String> {
keywords
.lines()
.filter(|keyword| !keyword.is_empty())
.map(str::to_string)
.collect()
}
pub(super) fn memory_token_estimate(
statement: &str,
relationships: &[ClaimRelationship],
match_reasons: &[MatchReason],
evidence: &[SourceReference],
) -> usize {
let relationship_chars: usize = relationships
.iter()
.map(|relationship| {
relationship.kind.as_str().len()
+ relationship.claim_id.len()
+ relationship.text.len()
+ relationship.rationale.len()
})
.sum();
let reason_chars: usize = match_reasons
.iter()
.map(|reason| reason.kind.as_str().len() + reason.detail.len())
.sum();
let source_chars: usize = evidence
.iter()
.map(|source| {
source.citation_id.len()
+ source.provider.as_str().len()
+ source.session_id.len()
+ source.entry_id.len()
+ source.role.len()
+ source.project.as_os_str().len()
+ source.source_path.as_os_str().len()
+ source.content_hash.len()
+ source.snippet.len()
})
.sum();
(statement.chars().count() + relationship_chars + reason_chars + source_chars)
.div_ceil(4)
.max(1)
}
pub(super) fn session_tombstone_key(provider: &str, session_id: &str) -> String {
format!("{provider}\0{session_id}")
}
pub(super) fn memory_type_counts(conn: &Connection) -> anyhow::Result<MemoryTypeCounts> {
let mut counts = MemoryTypeCounts::default();
let mut stmt = conn.prepare("SELECT memory_type, count(*) FROM claims GROUP BY memory_type")?;
let rows = stmt.query_map([], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
})?;
for row in rows {
let (memory_type, count) = row?;
let count = u64::try_from(count)?;
match memory_type.as_str() {
"decision" => counts.decisions = count,
"fact" => counts.facts = count,
"preference" => counts.preferences = count,
"procedure" => counts.procedures = count,
"lesson" => counts.lessons = count,
_ => return Err(MemoryError::UnknownMemoryType(memory_type).into()),
}
}
Ok(counts)
}
pub(super) fn count(conn: &Connection, table: &str) -> anyhow::Result<u64> {
let sql = format!("SELECT count(*) FROM {table}");
let value: i64 = conn.query_row(&sql, [], |row| row.get(0))?;
Ok(u64::try_from(value)?)
}
pub(super) fn count_where(conn: &Connection, sql: &str, value: &str) -> anyhow::Result<u64> {
let count: i64 = conn.query_row(sql, params![value], |row| row.get(0))?;
Ok(u64::try_from(count)?)
}
pub(super) fn count_two(
conn: &Connection,
sql: &str,
left: &str,
right: &str,
) -> anyhow::Result<u64> {
let count: i64 = conn.query_row(sql, params![left, right], |row| row.get(0))?;
Ok(u64::try_from(count)?)
}
pub(super) fn now_millis() -> i64 {
let millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or(Duration::ZERO)
.as_millis();
i64::try_from(millis).unwrap_or(i64::MAX)
}