use anyhow::Result; 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 entity/fact extraction pipeline pub struct IngestWorker { pool: PgPool, vector_store: Arc, embeddings: Arc, pipeline: Arc, } impl IngestWorker { /// 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 -> entities/facts/edges via pipeline -> temporal storage pub async fn process_ingest( &self, project: &str, ingest_id: &str, records: Vec<(String, String)>, // (content, source) ) -> Result<()> { tracing::info!("Processing ingest: project={}, id={}, records={}", project, ingest_id, records.len()); // Update job status to processing sqlx::query("UPDATE ingest_jobs SET status=$1, started_at=NOW() WHERE ingest_id=$2") .bind("processing") .bind(ingest_id) .execute(&self.pool) .await?; let mut total_entities = 0; let mut total_edges = 0; let mut total_reviews = 0; // 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), }; // 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 ); // 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 } } } // Mark job complete sqlx::query("UPDATE ingest_jobs SET status=$1, completed_at=NOW() WHERE ingest_id=$2") .bind("done") .bind(ingest_id) .execute(&self.pool) .await?; tracing::info!( "Ingest completed: {} (entities={}, edges={}, reviews={})", ingest_id, total_entities, total_edges, total_reviews ); Ok(()) } /// Process a single chunk pub async fn process_chunk(&self, project: &str, query_id: &str, content: &str, source: &str) -> Result<()> { let embedding = self.embeddings.embed_one(content).await?; let chunk = ChunkL0 { id: Uuid::new_v4(), project: project.to_string(), query_id: query_id.to_string(), source: source.to_string(), content: content.to_string(), tokens: (content.len() / 4) as i32, }; self.vector_store.store_chunk_l0(&chunk).await?; 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 PostgreSQL timestamp format let t_created_str = entity.t_created.to_string(); sqlx::query( "INSERT INTO memory_entity (id, project_id, name, entity_type, description, t_created, t_updated, confidence) VALUES ($1, $2, $3, $4, $5, $6::TIMESTAMPTZ, $7::TIMESTAMPTZ, $8) ON CONFLICT (id) DO NOTHING" ) .bind(&entity.id) .bind(&entity.project_id) .bind(&entity.name) .bind(entity.entity_type.as_str()) .bind(entity.summary.as_deref()) .bind(&t_created_str) .bind(&t_created_str) .bind(1.0_f32) // default confidence .execute(pool) .await?; Ok(()) } /// Save edge to database via raw SQL (normally would use EdgeRepo trait) /// NOTE: Production DB may have old schema. Gracefully skip if temporal columns missing. async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> { // Try temporal schema first (id, project_id, source_entity_id, etc) let result = sqlx::query( "INSERT INTO memory_edge (id, project_id, source_entity_id, target_entity_id, relation_type, fact, t_valid, t_invalid, t_created, confidence) VALUES ($1, $2, $3, $4, $5, $6, $7::TIMESTAMPTZ, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10) 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(edge.t_valid.map(|t| t.to_string())) .bind(edge.t_invalid.map(|t| t.to_string())) .bind(edge.t_created.to_string()) .bind(edge.confidence) .execute(pool) .await; match result { Ok(_) => Ok(()), Err(e) => { tracing::debug!("Temporal edge schema not available: {}. Skipping edge save (will be available after schema migration).", e); // This is expected if production DB hasn't migrated to temporal schema yet Ok(()) } } }