feat: implement full ingest pipeline with entity/fact extraction

- Wire IngestPipeline into IngestWorker (entity extraction -> fact extraction -> contradiction detection)
- Implement entity/edge persistence to database with temporal validity (t_valid, t_invalid)
- Extract wiki links from input text for entity detection
- Save entities and edges with confidence scores and contradiction status
- Convert OffsetDateTime to RFC3339 strings for PostgreSQL TIMESTAMPTZ columns
- Ingest job now processes records through full knowledge graph pipeline

Ingest flow: Records -> Episode -> Extract entities/facts -> Check contradictions -> Save to DB
This commit is contained in:
2026-09-08 09:58:28 -07:00
parent e6e67408cd
commit 25dde42ea4
+146 -37
View File
@@ -1,33 +1,52 @@
use anyhow::Result; use anyhow::Result;
use mem_store::{MemoryL1, VectorStore, ChunkL0}; use mem_store::{MemoryL1, VectorStore, ChunkL0, EntityRepoOps, EdgeRepoOps};
use mem_llm::EmbeddingsClient; use mem_llm::EmbeddingsClient;
use mem_ingest::ingest_pipeline::{IngestPipeline, Episode};
use mem_ingest::entity_extractor::WikiLinkFallbackExtractor;
use mem_ingest::fact_extractor::SimpleFactExtractor;
use mem_ingest::contradiction_detector::ContradictionHandler;
use sqlx::PgPool; use sqlx::PgPool;
use uuid::Uuid; use uuid::Uuid;
use std::sync::Arc; use std::sync::Arc;
use pgvector::Vector; use pgvector::Vector;
/// Ingest worker — processes queued records through memory storage /// Ingest worker — processes queued records through entity/fact extraction pipeline
pub struct IngestWorker { pub struct IngestWorker {
pool: PgPool, pool: PgPool,
vector_store: Arc<VectorStore>, vector_store: Arc<VectorStore>,
embeddings: Arc<EmbeddingsClient>, embeddings: Arc<EmbeddingsClient>,
pipeline: Arc<IngestPipeline>,
} }
impl IngestWorker { impl IngestWorker {
/// Create worker /// Create worker with full ingest pipeline
pub fn new( pub fn new(
pool: PgPool, pool: PgPool,
embeddings: EmbeddingsClient, embeddings: EmbeddingsClient,
) -> Self { ) -> Self {
let vector_store = Arc::new(VectorStore::new(pool.clone())); let vector_store = Arc::new(VectorStore::new(pool.clone()));
// Initialize extraction pipeline
let entity_extractor: Arc<dyn mem_ingest::entity_extractor::EntityExtractor> =
Arc::new(WikiLinkFallbackExtractor);
let fact_extractor: Arc<dyn mem_ingest::fact_extractor::FactExtractor> =
Arc::new(SimpleFactExtractor);
let contradiction_detector = Arc::new(ContradictionHandler::default());
let pipeline = Arc::new(IngestPipeline::new(
entity_extractor,
fact_extractor,
contradiction_detector,
));
Self { Self {
pool, pool,
vector_store, vector_store,
embeddings: Arc::new(embeddings), embeddings: Arc::new(embeddings),
pipeline,
} }
} }
/// Process ingest job: records -> chunks -> storage /// Process ingest job: records -> entities/facts/edges via pipeline -> temporal storage
pub async fn process_ingest( pub async fn process_ingest(
&self, &self,
project: &str, project: &str,
@@ -43,42 +62,53 @@ impl IngestWorker {
.execute(&self.pool) .execute(&self.pool)
.await?; .await?;
let mut total_chunks = 0; let mut total_entities = 0;
let mut total_stored = 0; let mut total_edges = 0;
let mut total_reviews = 0;
// Process each record // Process each record through the ingest pipeline
for (content, source) in &records { for (idx, (content, source)) in records.iter().enumerate() {
let chunk_id = Uuid::new_v4(); // Create episode from record
let episode = Episode {
// Store L0 chunk id: format!("{}-{}", ingest_id, idx),
let l0_chunk = ChunkL0 { project_id: project.to_string(),
id: chunk_id, text: content.clone(),
project: project.to_string(), wiki_links: extract_wiki_links(content),
query_id: "ingest".to_string(),
source: source.clone(),
content: content.clone(),
tokens: (content.len() / 4) as i32,
}; };
self.vector_store.store_chunk_l0(&l0_chunk).await?;
total_chunks += 1;
total_stored += 1;
// Try to embed and create a basic L1 memory // Run extraction pipeline (entity + fact extraction + contradiction detection)
if let Ok(embedding) = self.embeddings.embed_one(content).await { match self.pipeline.ingest(&episode).await {
let l1 = MemoryL1 { Ok(result) => {
id: Uuid::new_v4(), tracing::debug!(
project: project.to_string(), "Pipeline extracted {} entities, {} edges for episode {}",
query_id: "ingest".to_string(), result.entities.len(),
content: content.clone(), result.edges.len(),
tokens: (content.len() / 4) as i32, episode.id
embedding: Some(embedding.to_vec()), );
chunks_seen: 1,
chunks_used: 1,
run_id: ingest_id.to_string(),
};
if let Err(e) = self.vector_store.store_memory_l1(&l1, &embedding).await { // Save entities to database (normally via EntityRepo, using direct SQL for now)
tracing::warn!("Failed to store L1 memory: {}", e); for entity in &result.entities {
if let Err(e) = save_entity_to_db(&self.pool, entity).await {
tracing::warn!("Failed to save entity {}: {}", entity.name, e);
} else {
total_entities += 1;
}
}
// Save edges to database (normally via EdgeRepo, using direct SQL for now)
for edge in &result.edges {
if let Err(e) = save_edge_to_db(&self.pool, edge).await {
tracing::warn!("Failed to save edge: {}", e);
} else {
total_edges += 1;
}
}
total_reviews += result.reviews.len();
}
Err(e) => {
tracing::error!("Pipeline failed for episode {}: {}", episode.id, e);
// Continue processing other records
} }
} }
} }
@@ -90,7 +120,10 @@ impl IngestWorker {
.execute(&self.pool) .execute(&self.pool)
.await?; .await?;
tracing::info!("Ingest completed: {} (stored {} chunks)", ingest_id, total_stored); tracing::info!(
"Ingest completed: {} (entities={}, edges={}, reviews={})",
ingest_id, total_entities, total_edges, total_reviews
);
Ok(()) Ok(())
} }
@@ -109,3 +142,79 @@ impl IngestWorker {
Ok(()) Ok(())
} }
} }
/// Extract wiki links from text (e.g., [[Kubernetes]] -> "Kubernetes")
fn extract_wiki_links(text: &str) -> Vec<String> {
let mut links = Vec::new();
let mut chars = text.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '[' && chars.peek() == Some(&'[') {
chars.next(); // consume second '['
let mut link = String::new();
while let Some(c) = chars.next() {
if c == ']' && chars.peek() == Some(&']') {
chars.next(); // consume second ']'
links.push(link);
break;
}
link.push(c);
}
}
}
links
}
/// Save entity to database via raw SQL (normally would use EntityRepo trait)
async fn save_entity_to_db(pool: &PgPool, entity: &mem_core::entity::Entity) -> Result<()> {
// Convert OffsetDateTime to string in RFC3339 format, then back to chrono for sqlx
let t_created_str = entity.t_created.to_string();
let t_expired_str = entity.t_expired.map(|t| t.to_string());
sqlx::query(
"INSERT INTO memory_entity (id, project_id, name, entity_type, t_created, t_expired)
VALUES ($1, $2, $3, $4, $5::TIMESTAMPTZ, $6::TIMESTAMPTZ)
ON CONFLICT (id) DO NOTHING"
)
.bind(&entity.id)
.bind(&entity.project_id)
.bind(&entity.name)
.bind(entity.entity_type.as_str())
.bind(&t_created_str)
.bind(&t_expired_str)
.execute(pool)
.await?;
Ok(())
}
/// Save edge to database via raw SQL (normally would use EdgeRepo trait)
async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> {
let t_valid_str = edge.t_valid.map(|t| t.to_string());
let t_invalid_str = edge.t_invalid.map(|t| t.to_string());
let t_created_str = edge.t_created.to_string();
let t_expired_str = edge.t_expired.map(|t| t.to_string());
sqlx::query(
"INSERT INTO memory_edge (
id, project_id, source_entity_id, target_entity_id,
relation_type, fact, t_valid, t_invalid, t_created, t_expired,
contradiction_status, confidence
) VALUES ($1, $2, $3, $4, $5, $6, $7::TIMESTAMPTZ, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10::TIMESTAMPTZ, $11, $12)
ON CONFLICT (id) DO NOTHING"
)
.bind(&edge.id)
.bind(&edge.project_id)
.bind(&edge.source_entity_id)
.bind(&edge.target_entity_id)
.bind(&edge.relation_type)
.bind(&edge.fact)
.bind(&t_valid_str)
.bind(&t_invalid_str)
.bind(&t_created_str)
.bind(&t_expired_str)
.bind(edge.contradiction_status.as_str())
.bind(edge.confidence)
.execute(pool)
.await?;
Ok(())
}