diff --git a/crates/mem-cli/src/ingest_worker.rs b/crates/mem-cli/src/ingest_worker.rs index 732640e..b1317b7 100644 --- a/crates/mem-cli/src/ingest_worker.rs +++ b/crates/mem-cli/src/ingest_worker.rs @@ -1,33 +1,52 @@ use anyhow::Result; -use mem_store::{MemoryL1, VectorStore, ChunkL0}; +use mem_store::{MemoryL1, VectorStore, ChunkL0, EntityRepoOps, EdgeRepoOps}; 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 uuid::Uuid; use std::sync::Arc; use pgvector::Vector; -/// Ingest worker — processes queued records through memory storage +/// Ingest worker — processes queued records through entity/fact extraction pipeline pub struct IngestWorker { pool: PgPool, vector_store: Arc, embeddings: Arc, + pipeline: Arc, } impl IngestWorker { - /// Create worker + /// Create worker with full ingest pipeline pub fn new( pool: PgPool, embeddings: EmbeddingsClient, ) -> Self { let vector_store = Arc::new(VectorStore::new(pool.clone())); + + // Initialize extraction pipeline + let entity_extractor: Arc = + Arc::new(WikiLinkFallbackExtractor); + let fact_extractor: Arc = + Arc::new(SimpleFactExtractor); + let contradiction_detector = Arc::new(ContradictionHandler::default()); + let pipeline = Arc::new(IngestPipeline::new( + entity_extractor, + fact_extractor, + contradiction_detector, + )); + Self { pool, vector_store, 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( &self, project: &str, @@ -43,42 +62,53 @@ impl IngestWorker { .execute(&self.pool) .await?; - let mut total_chunks = 0; - let mut total_stored = 0; + let mut total_entities = 0; + let mut total_edges = 0; + let mut total_reviews = 0; - // Process each record - for (content, source) in &records { - let chunk_id = Uuid::new_v4(); - - // Store L0 chunk - let l0_chunk = ChunkL0 { - id: chunk_id, - project: project.to_string(), - query_id: "ingest".to_string(), - source: source.clone(), - content: content.clone(), - tokens: (content.len() / 4) as i32, + // Process each record through the ingest pipeline + for (idx, (content, source)) in records.iter().enumerate() { + // Create episode from record + let episode = Episode { + id: format!("{}-{}", ingest_id, idx), + project_id: project.to_string(), + text: content.clone(), + wiki_links: extract_wiki_links(content), }; - self.vector_store.store_chunk_l0(&l0_chunk).await?; - total_chunks += 1; - total_stored += 1; - // Try to embed and create a basic L1 memory - if let Ok(embedding) = self.embeddings.embed_one(content).await { - let l1 = MemoryL1 { - id: Uuid::new_v4(), - project: project.to_string(), - query_id: "ingest".to_string(), - content: content.clone(), - tokens: (content.len() / 4) as i32, - embedding: Some(embedding.to_vec()), - chunks_seen: 1, - chunks_used: 1, - run_id: ingest_id.to_string(), - }; + // Run extraction pipeline (entity + fact extraction + contradiction detection) + match self.pipeline.ingest(&episode).await { + Ok(result) => { + tracing::debug!( + "Pipeline extracted {} entities, {} edges for episode {}", + result.entities.len(), + result.edges.len(), + episode.id + ); - if let Err(e) = self.vector_store.store_memory_l1(&l1, &embedding).await { - tracing::warn!("Failed to store L1 memory: {}", e); + // Save entities to database (normally via EntityRepo, using direct SQL for now) + 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) .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(()) } @@ -109,3 +142,79 @@ impl IngestWorker { Ok(()) } } + +/// Extract wiki links from text (e.g., [[Kubernetes]] -> "Kubernetes") +fn extract_wiki_links(text: &str) -> Vec { + 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(()) +}