feat: LLM entity + fact extraction pipeline (Zep paper alignment) (#48)
CI / CI (push) Successful in 12m9s
Deploy / Tag & Push Latest (push) Failing after 41s
DB Migration / Run Migrations (push) Failing after 18s

## Changes

### Entity Extraction
- Switch from WikiLinkFallbackExtractor to LlmEntityExtractor when LLM_ENDPOINT set
- `clean_llm_response()`: strips `<think>` tags, markdown fences, extracts JSON
- Handle array responses (Ollama returns `[...]` not `{entities: [...]}`)
- EntityType custom Deserialize: unknown variants → Unknown (no crash)
- Increase timeout 30s→90s, max_tokens 500→1500 for reasoning models
- Graceful reflection fallback: keep entities if verification fails

### Fact Extraction (NEW)
- LlmFactExtractor: LLM-based relationship extraction between entity pairs
- Validates source/target against known entity list (drops hallucinated edges)
- Same robust JSON cleaning for reasoning models + Ollama
- IngestWorker auto-selects LLM vs Simple based on LLM_ENDPOINT env

### K8s Deployment
- Add `command: ["/app/mem"]` (fix args replacing CMD)
- Add LLM_ENDPOINT, LLM_MODEL env vars for in-cluster LLM

## E2E Tested (local Ollama qwen2.5:3b)
- 12 entities extracted (person, tool, concept, organization)
- 5 edges with relationships and facts
- 781 tests pass

## Zep Paper Alignment (§2.2)
- Entity extraction + resolution (§2.2.1)
- Fact extraction between entity pairs (§2.2.2)
- Temporal edge invalidation ready (t_valid/t_invalid schema)
- Reflection verification (§2.2.1, graceful fallback)

---------

Co-authored-by: rock <[email protected]>
Reviewed-on: #48
Co-authored-by: poimen <[email protected]>
This commit was merged in pull request #48.
This commit is contained in:
2026-09-11 01:11:15 +00:00
committed by rock
co-authored by rock
parent 6b18d81421
commit fb61de6b47
19 changed files with 859 additions and 268 deletions
+35 -9
View File
@@ -2,14 +2,15 @@ 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::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;
use pgvector::Vector;
/// Ingest worker — processes queued records through entity/fact extraction pipeline
pub struct IngestWorker {
pool: PgPool,
@@ -26,11 +27,25 @@ impl IngestWorker {
) -> Self {
let vector_store = Arc::new(VectorStore::new(pool.clone()));
// Initialize extraction pipeline
// Initialize extraction pipeline — use LLM if LLM_ENDPOINT is set, else fallback to wiki links
let entity_extractor: Arc<dyn mem_ingest::entity_extractor::EntityExtractor> =
Arc::new(WikiLinkFallbackExtractor);
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<dyn mem_ingest::fact_extractor::FactExtractor> =
Arc::new(SimpleFactExtractor);
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,
@@ -121,9 +136,15 @@ impl IngestWorker {
.await?;
tracing::info!(
"Ingest completed: {} (entities={}, edges={}, reviews={})",
ingest_id, total_entities, total_edges, total_reviews
target: "observability",
event = "ingest_complete",
ingest_id = ingest_id,
entities = total_entities,
edges = total_edges,
reviews = total_reviews,
"Ingest completed"
);
Ok(())
}
@@ -173,7 +194,12 @@ async fn save_entity_to_db(pool: &PgPool, entity: &mem_core::entity::Entity) ->
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"
ON CONFLICT (project_id, name) DO UPDATE SET
entity_type = EXCLUDED.entity_type,
description = COALESCE(NULLIF(EXCLUDED.description, ''), memory_entity.description),
t_updated = NOW(),
confidence = GREATEST(memory_entity.confidence, EXCLUDED.confidence),
source_count = memory_entity.source_count + 1"
)
.bind(&entity.id)
.bind(&entity.project_id)
@@ -193,7 +219,7 @@ async fn save_entity_to_db(pool: &PgPool, entity: &mem_core::entity::Entity) ->
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)
"INSERT INTO memory_edge (id, project_id, source_id, target_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"
)