use anyhow::Result; use mem_store::{VectorStore, ChunkL0}; use mem_llm::EmbeddingsClient; use mem_ingest::ingest_pipeline::{IngestPipeline, Episode}; use mem_ingest::entity_extractor::{WikiLinkFallbackExtractor, LlmEntityExtractor}; use mem_ingest::fact_extractor::{SimpleFactExtractor, LlmFactExtractor}; use mem_ingest::contradiction_detector::ContradictionHandler; use sqlx::PgPool; use uuid::Uuid; use std::sync::Arc; /// Job status enumeration — type-safe alternative to magic strings #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[allow(dead_code)] pub enum JobStatus { Processing, Done, DoneWithErrors, } #[allow(dead_code)] impl JobStatus { pub fn as_str(&self) -> &'static str { match self { JobStatus::Processing => "processing", JobStatus::Done => "done", JobStatus::DoneWithErrors => "done_with_errors", } } } impl std::fmt::Display for JobStatus { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.as_str()) } } #[cfg(test)] mod tests { use super::*; /// Mock JobStatusStore for testing pub struct MockJobStatusStore { updates: std::sync::Arc>>, } impl MockJobStatusStore { pub fn new() -> Self { Self { updates: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())), } } pub fn updates(&self) -> Vec<(String, JobStatus)> { self.updates.lock().unwrap().clone() } } #[async_trait::async_trait] impl JobStatusStore for MockJobStatusStore { async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()> { self.updates.lock().unwrap().push((ingest_id.to_string(), status)); Ok(()) } } } /// Structured logging context for ingest operations — ensures consistent field names across all logs #[derive(Debug, Clone)] #[allow(dead_code)] pub struct IngestLogContext { pub ingest_id: String, pub project: String, pub record_id: String, pub source: String, } #[allow(dead_code)] impl IngestLogContext { fn new(ingest_id: &str, project: &str, record_id: &str, source: &str) -> Self { Self { ingest_id: ingest_id.to_string(), project: project.to_string(), record_id: record_id.to_string(), source: source.to_string(), } } } /// Job status store trait — abstracts database persistence of job status (enables mocking) #[async_trait::async_trait] #[allow(dead_code)] pub trait JobStatusStore: Send + Sync { /// Update job status in storage async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()>; } /// PostgreSQL implementation of JobStatusStore #[allow(dead_code)] pub struct PgJobStatusStore { pool: PgPool, } #[allow(dead_code)] impl PgJobStatusStore { pub fn new(pool: PgPool) -> Self { Self { pool } } } #[async_trait::async_trait] impl JobStatusStore for PgJobStatusStore { async fn update_status(&self, ingest_id: &str, status: JobStatus) -> Result<()> { sqlx::query("UPDATE ingest_jobs SET status=$1, started_at=NOW() WHERE ingest_id=$2") .bind(status.as_str()) .bind(ingest_id) .execute(&self.pool) .await?; Ok(()) } } /// Ingest worker — processes queued records through entity/fact extraction pipeline #[allow(dead_code)] pub struct IngestWorker { pool: PgPool, vector_store: Arc, embeddings: Arc, pipeline: Arc, job_status_store: Arc, } #[allow(dead_code)] impl IngestWorker { /// Create worker with full ingest pipeline pub fn new( pool: PgPool, embeddings: EmbeddingsClient, ) -> Self { let job_status_store = Arc::new(PgJobStatusStore::new(pool.clone())); Self::with_job_store(pool, embeddings, job_status_store) } /// Create worker with custom job status store (for testing) pub fn with_job_store( pool: PgPool, embeddings: EmbeddingsClient, job_status_store: Arc, ) -> Self { let vector_store = Arc::new(VectorStore::new(pool.clone())); // Initialize extraction pipeline — use LLM if LLM_ENDPOINT is set, else fallback to wiki links let entity_extractor: Arc = if std::env::var("LLM_ENDPOINT").is_ok() { let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "qwen2.5:3b-instruct".to_string()); tracing::info!("Using LLM entity extractor: model={}", model); Arc::new(LlmEntityExtractor::new(&model)) } else { tracing::info!("LLM_ENDPOINT not set, using WikiLink fallback extractor"); Arc::new(WikiLinkFallbackExtractor) }; let fact_extractor: Arc = if std::env::var("LLM_ENDPOINT").is_ok() { let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "qwen2.5:3b-instruct".to_string()); tracing::info!("Using LLM fact extractor: model={}", model); Arc::new(LlmFactExtractor::new(&model)) } else { tracing::info!("LLM_ENDPOINT not set, using simple pattern fact extractor"); 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, job_status_store, } } /// Process ingest job with optional X-Forward-User auth header (API Gateway pattern) /// /// # Arguments /// * `project` - Project ID for namespacing /// * `ingest_id` - Unique ingest job ID /// * `records` - Vec of (content, source) tuples /// * `x_forward_user` - Optional X-Forward-User header from API Gateway (None for backward compat) pub async fn process_ingest_with_auth( &self, project: &str, ingest_id: &str, records: Vec<(String, String)>, // (content, source) x_forward_user: Option, ) -> Result<()> { tracing::info!( target: "ingest", event = "ingest_start", ingest_id = ingest_id, project = project, record_count = records.len(), "Starting ingest job" ); // Update job status to processing (via trait, testable) if let Err(e) = self.job_status_store.update_status(ingest_id, JobStatus::Processing).await { tracing::error!( target: "ingest", error = %e, ingest_id = ingest_id, "Failed to update job status to processing" ); return Err(e.into()); } 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() { let record_id = format!("{}-{}", ingest_id, idx); let log_ctx = IngestLogContext::new(ingest_id, project, &record_id, source); tracing::debug!( target: "ingest", record_id = %log_ctx.record_id, source = %log_ctx.source, content_len = content.len(), "Processing record" ); // Create episode from record let episode = Episode { id: record_id.clone(), project_id: project.to_string(), text: content.clone(), wiki_links: extract_wiki_links(content), }; // Run extraction pipeline (entity + fact extraction + contradiction detection) let x_forward_user_ref = x_forward_user.as_deref(); match self.pipeline.ingest_with_auth(&episode, x_forward_user_ref).await { Ok(result) => { tracing::debug!( target: "ingest", record_id = %log_ctx.record_id, entity_count = result.entities.len(), edge_count = result.edges.len(), review_count = result.reviews.len(), "Pipeline extraction successful" ); // Save entities to database with embeddings (RAG-006) for entity in &result.entities { match save_entity_with_embedding(&self.pool, &self.embeddings, entity, &log_ctx).await { Ok(saved) => if saved { total_entities += 1; } Err(_) => { /* error already logged */ } } } // Save edges to database with embeddings (RAG-006) for edge in &result.edges { match save_edge_with_embedding(&self.pool, &self.embeddings, edge, &log_ctx).await { Ok(saved) => if saved { total_edges += 1; } Err(_) => { /* error already logged */ } } } total_reviews += result.reviews.len(); } Err(e) => { tracing::error!( target: "ingest", error = %e, record_id = %log_ctx.record_id, source = %log_ctx.source, "Pipeline extraction failed" ); // Continue processing other records (no error accumulation) } } } // Mark job complete (via trait, testable) let final_status = JobStatus::Done; if let Err(e) = self.job_status_store.update_status(ingest_id, final_status).await { tracing::error!( target: "ingest", error = %e, ingest_id = ingest_id, "Failed to update job completion status" ); } tracing::info!( target: "ingest", event = "ingest_complete", ingest_id = ingest_id, project = project, entities = total_entities, edges = total_edges, reviews = total_reviews, status = final_status.as_str(), "Ingest job completed" ); 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") #[allow(dead_code)] 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 with logging — logs at debug level on success, warn on error /// Returns Ok(true) if saved, Ok(false) if skipped, Err if fatal error /// Save entity with embeddings (RAG-006) /// Embeds name + summary before persisting, so semantic search can find entities. #[allow(dead_code)] async fn save_entity_with_embedding( pool: &PgPool, embeddings: &EmbeddingsClient, entity: &mem_core::entity::Entity, log_ctx: &IngestLogContext, ) -> Result { // Embed entity name let name_embedding = match embeddings.embed_one(&entity.name).await { Ok(emb) => Some(emb.to_vec()), Err(e) => { tracing::warn!( target: "ingest", error = %e, entity_name = &entity.name, "Name embedding failed, saving entity without name_embedding" ); None } }; // Embed summary if present let summary_embedding = if let Some(ref summary) = entity.summary { match embeddings.embed_one(summary).await { Ok(emb) => Some(emb.to_vec()), Err(e) => { tracing::debug!(target: "ingest", error = %e, "Summary embedding failed"); None } } } else { None }; let t_created_str = entity.t_created.to_string(); let result = sqlx::query( "INSERT INTO memory_entity (id, project_id, name, entity_type, description, summary, \ name_embedding, summary_embedding, t_created, t_updated, confidence) \ VALUES ($1::UUID, $2, $3, $4, $5, $6, $7, $8, $9::TIMESTAMPTZ, $10::TIMESTAMPTZ, $11) \ ON CONFLICT (project_id, name) DO UPDATE SET \ entity_type = EXCLUDED.entity_type, \ description = COALESCE(NULLIF(EXCLUDED.description, ''), memory_entity.description), \ summary = COALESCE(NULLIF(EXCLUDED.summary, ''), memory_entity.summary), \ name_embedding = COALESCE(EXCLUDED.name_embedding, memory_entity.name_embedding), \ summary_embedding = COALESCE(EXCLUDED.summary_embedding, memory_entity.summary_embedding), \ t_updated = NOW(), \ confidence = GREATEST(memory_entity.confidence, EXCLUDED.confidence), \ source_count = memory_entity.source_count + 1" ) .bind(&entity.id) .bind(&entity.project_id) .bind(&entity.name) .bind(entity.entity_type.as_str()) .bind(entity.summary.as_deref()) // description .bind(entity.summary.as_deref()) // summary .bind(name_embedding.as_deref()) .bind(summary_embedding.as_deref()) .bind(&t_created_str) .bind(&t_created_str) .bind(1.0_f32) .execute(pool) .await; match result { Ok(_) => { tracing::debug!( target: "ingest", record_id = %log_ctx.record_id, entity_name = &entity.name, entity_type = entity.entity_type.as_str(), has_name_emb = name_embedding.is_some(), has_summary_emb = summary_embedding.is_some(), "Saved entity with embeddings" ); Ok(true) } Err(e) => { tracing::warn!( target: "ingest", error = %e, record_id = %log_ctx.record_id, entity_name = &entity.name, project = %log_ctx.project, "Entity save failed" ); Ok(false) } } } /// Save edge with fact embedding (RAG-006) /// Embeds fact text before persisting, so semantic search can find edges. #[allow(dead_code)] async fn save_edge_with_embedding( pool: &PgPool, embeddings: &EmbeddingsClient, edge: &mem_core::edge::Edge, log_ctx: &IngestLogContext, ) -> Result { // Embed the fact text let fact_embedding = match embeddings.embed_one(&edge.fact).await { Ok(emb) => Some(emb.to_vec()), Err(e) => { tracing::warn!( target: "ingest", error = %e, fact = &edge.fact, "Fact embedding failed, saving edge without fact_embedding" ); None } }; let result = sqlx::query( "INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, \ fact_embedding, t_valid, t_invalid, t_created, confidence) \ VALUES ($1::UUID, $2, $3::UUID, $4::UUID, $5, $6, $7, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10::TIMESTAMPTZ, $11) \ 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(fact_embedding.as_deref()) .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(_) => { tracing::debug!( target: "ingest", record_id = %log_ctx.record_id, relation_type = &edge.relation_type, source_entity = &edge.source_entity_id, target_entity = &edge.target_entity_id, has_fact_emb = fact_embedding.is_some(), "Saved edge with embedding" ); Ok(true) } Err(e) => { tracing::warn!( target: "ingest", error = %e, record_id = %log_ctx.record_id, relation_type = &edge.relation_type, project = %log_ctx.project, "Edge save failed" ); Ok(false) } } } // Legacy save functions kept for backward compatibility but unused #[allow(dead_code)] #[allow(dead_code)] async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> { let result = sqlx::query( "INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, t_valid, t_invalid, t_created, confidence) VALUES ($1::UUID, $2, $3::UUID, $4::UUID, $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.", e); Ok(()) } } }