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:
@@ -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(())
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user