use std::collections::BTreeSet;
use anyhow::{Context, Result};
use rusqlite::{params, Connection, OptionalExtension};
use crate::db;
use crate::graph_candidate::ParsedGraphCandidate;
use crate::memory::format::{xml_escape_attr, xml_escape_text};
#[derive(Debug, Clone)]
struct GraphSourceObservation {
id: i64,
observation_type: String,
text: String,
files_read: Vec<String>,
files_modified: Vec<String>,
evidence_event_ids: Vec<i64>,
confidence: f64,
}
#[derive(Debug, Clone)]
struct GraphSourceEvent {
id: i64,
text: String,
}
pub(super) struct GraphObservationBatch {
pub(super) from_event_id: i64,
pub(super) to_event_id: i64,
pub(super) evidence_event_ids: Vec<i64>,
source_events: Vec<GraphSourceEvent>,
observations: Vec<GraphSourceObservation>,
}
pub(super) fn load_graph_observation_batch(
conn: &Connection,
task: &db::ExtractionTask,
) -> Result<Option<GraphObservationBatch>> {
let Some(session_row_id) = task.session_row_id else {
return Ok(None);
};
let Some(high_watermark) = task.high_watermark_event_id else {
return Ok(None);
};
let cursor = task.cursor_event_id.unwrap_or(0);
if high_watermark <= cursor {
return Ok(None);
}
let mut stmt = conn.prepare(
"SELECT id,
COALESCE(observation_type, type, 'discovery') AS observation_type,
COALESCE(text, narrative, title, '') AS text,
files_read,
files_modified,
evidence_event_ids,
COALESCE(confidence, 0.5) AS confidence
FROM observations
WHERE session_row_id = ?1
AND evidence_event_ids IS NOT NULL
AND text IS NOT NULL
ORDER BY id ASC",
)?;
let rows = stmt
.query_map(params![session_row_id], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, Option<String>>(3)?,
row.get::<_, Option<String>>(4)?,
row.get::<_, String>(5)?,
row.get::<_, f64>(6)?,
))
})?
.collect::<Result<Vec<_>, _>>()?;
let mut observations = Vec::new();
let mut evidence_set = BTreeSet::new();
for (
id,
observation_type,
text,
files_read_json,
files_modified_json,
evidence_json,
confidence,
) in rows
{
let event_ids: Vec<i64> = serde_json::from_str(&evidence_json)
.with_context(|| format!("observation {id} has malformed evidence_event_ids"))?;
let in_range = event_ids
.iter()
.any(|event_id| *event_id > cursor && *event_id <= high_watermark);
if !in_range {
continue;
}
for event_id in &event_ids {
evidence_set.insert(*event_id);
}
observations.push(GraphSourceObservation {
id,
observation_type,
text,
files_read: parse_observation_file_list(id, "files_read", files_read_json)?,
files_modified: parse_observation_file_list(id, "files_modified", files_modified_json)?,
evidence_event_ids: event_ids,
confidence,
});
}
if observations.is_empty() || evidence_set.is_empty() {
return Ok(None);
}
let from_event_id = *evidence_set.iter().next().unwrap_or(&0);
let to_event_id = *evidence_set.iter().next_back().unwrap_or(&0);
let evidence_event_ids = evidence_set.into_iter().collect::<Vec<_>>();
let source_events = load_graph_source_events(conn, &evidence_event_ids)?;
Ok(Some(GraphObservationBatch {
from_event_id,
to_event_id,
evidence_event_ids,
source_events,
observations,
}))
}
fn load_graph_source_events(
conn: &Connection,
evidence_event_ids: &[i64],
) -> Result<Vec<GraphSourceEvent>> {
let mut stmt = conn.prepare(
"SELECT COALESCE(
CASE
WHEN b.content_encoding = 'plain' THEN CAST(b.content_bytes AS TEXT)
ELSE NULL
END,
e.content_text,
''
) AS content
FROM captured_events e
LEFT JOIN event_blobs b ON b.id = e.content_blob_id
WHERE e.id = ?1",
)?;
let mut events = Vec::new();
for event_id in evidence_event_ids {
if let Some(text) = stmt
.query_row(params![event_id], |row| row.get(0))
.optional()?
{
events.push(GraphSourceEvent {
id: *event_id,
text,
});
}
}
Ok(events)
}
fn parse_observation_file_list(
observation_id: i64,
field: &str,
raw: Option<String>,
) -> Result<Vec<String>> {
let Some(raw) = raw.as_deref().map(str::trim).filter(|raw| !raw.is_empty()) else {
return Ok(Vec::new());
};
serde_json::from_str::<Vec<String>>(raw)
.with_context(|| format!("observation {observation_id} has malformed {field}"))
}
pub(super) fn graph_candidate_blocked_by_memory_candidates(
conn: &Connection,
task: &db::ExtractionTask,
batch: &GraphObservationBatch,
) -> Result<Option<String>> {
let Some(session_row_id) = task.session_row_id else {
return Ok(None);
};
let high_watermark = task.high_watermark_event_id.unwrap_or(batch.to_event_id);
let incomplete_memory_task: Option<(i64, String)> = conn
.query_row(
"SELECT id, status
FROM extraction_tasks
WHERE host_id = ?1
AND project_id = ?2
AND session_row_id = ?3
AND task_kind = 'memory_candidate'
AND COALESCE(high_watermark_event_id, 0) >= ?4
AND COALESCE(cursor_event_id, 0) < ?4
AND status != 'done'
ORDER BY id ASC
LIMIT 1",
params![
task.host_id,
task.project_id,
session_row_id,
high_watermark
],
|row| Ok((row.get(0)?, row.get(1)?)),
)
.optional()?;
if let Some((task_id, status)) = incomplete_memory_task {
return Ok(Some(format!(
"memory_candidate task {task_id} is {status}; graph extraction waits for memory extraction completion"
)));
}
let pending_memory_candidates = count_pending_memory_candidates_with_evidence_overlap(
conn,
task.project_id,
&batch.evidence_event_ids,
)?;
if pending_memory_candidates > 0 {
return Ok(Some(format!(
"{pending_memory_candidates} memory candidate(s) for this evidence range are pending review"
)));
}
Ok(None)
}
fn count_pending_memory_candidates_with_evidence_overlap(
conn: &Connection,
project_id: i64,
evidence_event_ids: &[i64],
) -> Result<i64> {
let evidence_event_ids = evidence_event_ids.iter().copied().collect::<BTreeSet<_>>();
let mut stmt = conn.prepare(
"SELECT id, evidence_event_ids
FROM memory_candidates
WHERE project_id = ?1
AND review_status = 'pending_review'
AND evidence_event_ids IS NOT NULL",
)?;
let rows = stmt
.query_map(params![project_id], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?))
})?
.collect::<Result<Vec<_>, _>>()?;
let mut count = 0;
for (candidate_id, candidate_evidence_json) in rows {
let candidate_event_ids = serde_json::from_str::<Vec<i64>>(&candidate_evidence_json)
.with_context(|| {
format!("memory candidate {candidate_id} has malformed evidence_event_ids")
})?;
if candidate_event_ids
.iter()
.any(|event_id| evidence_event_ids.contains(event_id))
{
count += 1;
}
}
Ok(count)
}
pub(super) fn graph_candidate_has_source_support(
candidate: &ParsedGraphCandidate,
batch: &GraphObservationBatch,
) -> bool {
let candidate_evidence = candidate
.evidence_event_ids
.iter()
.copied()
.collect::<BTreeSet<_>>();
let observation_supported = batch
.observations
.iter()
.filter(|observation| {
observation
.evidence_event_ids
.iter()
.any(|event_id| candidate_evidence.contains(event_id))
})
.any(|observation| refs_supported_by_observation(candidate, observation));
let cited_events_supported = candidate.evidence_event_ids.iter().all(|event_id| {
let event_text = batch
.source_events
.iter()
.find(|event| event.id == *event_id)
.map(|event| event.text.as_str());
let event_text_supported =
event_text.is_some_and(|text| refs_supported_by_event(candidate, *event_id, text));
event_text_supported
|| batch.observations.iter().any(|observation| {
observation.evidence_event_ids.contains(event_id)
&& structured_file_evidence_supports_cited_event(
candidate,
*event_id,
event_text,
observation,
)
})
});
observation_supported && cited_events_supported
}
fn refs_supported_by_observation(
candidate: &ParsedGraphCandidate,
observation: &GraphSourceObservation,
) -> bool {
ref_supported_by_observation(&candidate.from_ref, observation)
&& ref_supported_by_observation(&candidate.to_ref, observation)
}
fn ref_supported_by_observation(reference: &str, observation: &GraphSourceObservation) -> bool {
parse_episode_ref_id(reference)
.is_some_and(|episode_id| observation.evidence_event_ids.contains(&episode_id))
|| ref_supported_by_text(reference, &observation_support_text(observation))
}
fn structured_file_evidence_supports_cited_event(
candidate: &ParsedGraphCandidate,
event_id: i64,
event_text: Option<&str>,
observation: &GraphSourceObservation,
) -> bool {
candidate.edge_type == "touches_file"
&& event_text
.is_some_and(|text| ref_supported_by_event(&candidate.from_ref, event_id, text))
&& ref_supported_by_structured_files(&candidate.to_ref, observation)
}
fn ref_supported_by_structured_files(
reference: &str,
observation: &GraphSourceObservation,
) -> bool {
observation
.files_read
.iter()
.chain(observation.files_modified.iter())
.any(|file| ref_supported_by_text(reference, file))
}
fn observation_support_text(observation: &GraphSourceObservation) -> String {
let mut text = observation.text.clone();
for file in observation
.files_read
.iter()
.chain(observation.files_modified.iter())
{
text.push('\n');
text.push_str(file);
}
text
}
fn refs_supported_by_event(candidate: &ParsedGraphCandidate, event_id: i64, text: &str) -> bool {
ref_supported_by_event(&candidate.from_ref, event_id, text)
&& ref_supported_by_event(&candidate.to_ref, event_id, text)
}
fn ref_supported_by_text(reference: &str, text: &str) -> bool {
let haystack = text.to_ascii_lowercase();
reference_support_needles(reference)
.into_iter()
.any(|needle| contains_support_needle(&haystack, &needle))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SupportBoundary {
Word,
Entity,
Path,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct SupportNeedle {
text: String,
boundary: SupportBoundary,
}
fn reference_support_needles(reference: &str) -> Vec<SupportNeedle> {
let reference = reference.trim();
if let Some(path) = reference.strip_prefix("file:") {
return vec![
SupportNeedle {
text: path.trim_start_matches("./").to_string(),
boundary: SupportBoundary::Path,
},
SupportNeedle {
text: path.to_string(),
boundary: SupportBoundary::Path,
},
];
}
if let Some(entity) = reference.strip_prefix("entity:") {
return vec![
SupportNeedle {
text: entity.to_string(),
boundary: SupportBoundary::Entity,
},
SupportNeedle {
text: entity.replace(['_', '-'], " "),
boundary: SupportBoundary::Entity,
},
];
}
if let Some(memory_id) = reference.strip_prefix("memory:") {
return vec![
SupportNeedle {
text: reference.to_string(),
boundary: SupportBoundary::Word,
},
SupportNeedle {
text: format!("memory {memory_id}"),
boundary: SupportBoundary::Word,
},
];
}
if let Some(episode_id) = reference.strip_prefix("episode:") {
return vec![
SupportNeedle {
text: reference.to_string(),
boundary: SupportBoundary::Word,
},
SupportNeedle {
text: format!("event {episode_id}"),
boundary: SupportBoundary::Word,
},
];
}
vec![SupportNeedle {
text: reference.to_string(),
boundary: SupportBoundary::Word,
}]
}
fn contains_support_needle(haystack: &str, needle: &SupportNeedle) -> bool {
let needle_text = needle.text.trim().to_ascii_lowercase();
!needle_text.is_empty()
&& haystack.match_indices(&needle_text).any(|(start, _)| {
let end = start + needle_text.len();
support_boundary_before(haystack, start, needle.boundary)
&& support_boundary_after(haystack, end, needle.boundary)
})
}
fn support_boundary_before(haystack: &str, start: usize, boundary: SupportBoundary) -> bool {
haystack[..start]
.chars()
.next_back()
.is_none_or(|ch| is_support_boundary(ch, boundary))
}
fn support_boundary_after(haystack: &str, end: usize, boundary: SupportBoundary) -> bool {
let mut chars = haystack[end..].chars();
let Some(ch) = chars.next() else {
return true;
};
if boundary == SupportBoundary::Path
&& ch == '.'
&& chars
.next()
.is_some_and(|next| next.is_ascii_alphanumeric())
{
return false;
}
is_support_boundary(ch, boundary)
}
fn is_support_boundary(ch: char, boundary: SupportBoundary) -> bool {
if ch.is_ascii_alphanumeric() || ch == '_' {
return false;
}
match boundary {
SupportBoundary::Word => true,
SupportBoundary::Entity => !matches!(ch, '-' | '/' | '.'),
SupportBoundary::Path => !matches!(ch, '-' | '/'),
}
}
fn ref_supported_by_event(reference: &str, event_id: i64, text: &str) -> bool {
parse_episode_ref_id(reference).is_some_and(|episode_id| episode_id == event_id)
|| ref_supported_by_text(reference, text)
}
fn parse_episode_ref_id(reference: &str) -> Option<i64> {
reference
.strip_prefix("episode:")
.and_then(|raw| raw.trim().parse::<i64>().ok())
.filter(|id| *id > 0)
}
pub(super) fn build_graph_candidate_prompt(
task: &db::ExtractionTask,
batch: &GraphObservationBatch,
) -> String {
let mut prompt = format!(
"Task: graph_candidate\nProject: {}\nHost: {}\nSession: {}\nCovered evidence events: {}..{}\n\n",
task.project,
task.host,
task.session_id.as_deref().unwrap_or("<unknown>"),
batch.from_event_id,
batch.to_event_id
);
for observation in &batch.observations {
let evidence = observation
.evidence_event_ids
.iter()
.map(i64::to_string)
.collect::<Vec<_>>()
.join(",");
prompt.push_str(&format!(
"<observation id=\"{}\" type=\"{}\" confidence=\"{}\" evidence_event_ids=\"{}\">\n",
observation.id,
xml_escape_attr(&observation.observation_type),
observation.confidence,
xml_escape_attr(&evidence)
));
prompt.push_str(&xml_escape_text(&observation.text));
if !observation.files_read.is_empty() {
prompt.push_str("\n<files_read>\n");
for file in &observation.files_read {
prompt.push_str(&xml_escape_text(file));
prompt.push('\n');
}
prompt.push_str("</files_read>\n");
}
if !observation.files_modified.is_empty() {
prompt.push_str("\n<files_modified>\n");
for file in &observation.files_modified {
prompt.push_str(&xml_escape_text(file));
prompt.push('\n');
}
prompt.push_str("</files_modified>\n");
}
prompt.push_str("\n</observation>\n\n");
}
prompt
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn support_matching_uses_ref_boundaries() {
assert!(!ref_supported_by_text(
"memory:1",
"Memory 10 mentions Worker."
));
assert!(ref_supported_by_text(
"memory:1",
"Memory 1 mentions Worker."
));
assert!(!ref_supported_by_text(
"entity:Worker",
"Changed src/worker.rs"
));
assert!(ref_supported_by_text(
"entity:Worker",
"The Worker entity is mentioned."
));
assert!(!ref_supported_by_text(
"file:src/worker.rs",
"Changed src/worker.rs.bak"
));
assert!(ref_supported_by_text(
"file:src/worker.rs",
"Changed src/worker.rs."
));
}
}