poimen-memory/crates/mem-cli/src/ingest_worker.rs

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<VectorStore>,
    embeddings: Arc<EmbeddingsClient>,
    pipeline: Arc<IngestPipeline>,
}

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<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 {
            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<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 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(())
        }
    }
}