diff --git a/crates/mem-cli/src/ingest_worker.rs b/crates/mem-cli/src/ingest_worker.rs index 4a10a0c..6b6c910 100644 --- a/crates/mem-cli/src/ingest_worker.rs +++ b/crates/mem-cli/src/ingest_worker.rs @@ -2,7 +2,7 @@ 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::entity_extractor::{WikiLinkFallbackExtractor, LlmEntityExtractor}; use mem_ingest::fact_extractor::SimpleFactExtractor; use mem_ingest::contradiction_detector::ContradictionHandler; use sqlx::PgPool; @@ -26,9 +26,16 @@ 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 = - 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 = Arc::new(SimpleFactExtractor); let contradiction_detector = Arc::new(ContradictionHandler::default()); diff --git a/crates/mem-ingest/src/entity_extractor.rs b/crates/mem-ingest/src/entity_extractor.rs index 60d533d..3a17e6a 100644 --- a/crates/mem-ingest/src/entity_extractor.rs +++ b/crates/mem-ingest/src/entity_extractor.rs @@ -22,11 +22,15 @@ use tokio::sync::Mutex; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ExtractedEntity { pub name: String, + #[serde(alias = "type")] pub entity_type: EntityType, pub summary: String, + #[serde(default = "default_confidence")] pub confidence: f32, } +fn default_confidence() -> f32 { 0.8 } + impl ExtractedEntity { /// Convert to domain model (Phase 1 type) pub fn to_domain(&self, project_id: &str) -> Entity { @@ -62,6 +66,25 @@ impl LlmEntityExtractor { /// Parse extraction response JSON /// Format: { "entities": [{ "name": "...", "type": "...", "summary": "..." }, ...] } + /// Strip ... tags from reasoning model output and extract JSON + fn strip_thinking_tags(text: &str) -> String { + let mut result = text.to_string(); + // Remove ... blocks + if let Some(start) = result.find("") { + if let Some(end) = result.find("") { + result = format!("{}{}", &result[..start], &result[end + 8..]); + } + } + // Try to find JSON object in remaining text + let trimmed = result.trim(); + if let Some(start) = trimmed.find('{') { + if let Some(end) = trimmed.rfind('}') { + return trimmed[start..=end].to_string(); + } + } + trimmed.to_string() + } + fn parse_extraction(response: &str) -> Result> { #[derive(Deserialize)] struct Response { @@ -146,12 +169,16 @@ impl LlmEntityExtractor { } let data: serde_json::Value = response.json().await?; - let content = data["choices"][0]["message"]["content"] + let raw_content = data["choices"][0]["message"]["content"] .as_str() .unwrap_or("{}") .to_string(); - tracing::debug!("LLM response (via Authentik JWT): {}", content); + // Strip ... tags from reasoning models + let content = Self::strip_thinking_tags(&raw_content); + + tracing::debug!("LLM raw response length={}, cleaned length={}", raw_content.len(), content.len()); + tracing::debug!("LLM cleaned content: {}", content); Ok(content) } @@ -233,14 +260,27 @@ Respond in JSON: ); let reflection = if std::env::var("LLM_ENDPOINT").is_ok() { - self.call_llm_endpoint(&reflection_prompt).await.unwrap_or_else(|_| self.simulate_llm(&reflection_prompt).unwrap_or_default()) + self.call_llm_endpoint(&reflection_prompt).await.unwrap_or_else(|e| { + tracing::warn!("Reflection LLM call failed: {}, skipping verification", e); + String::new() + }) } else { self.simulate_llm(&reflection_prompt)? }; - let verified = Self::parse_reflection(&reflection)?; - // Filter: keep only entities marked present - entities.retain(|e| verified.iter().any(|(name, present)| name == &e.name && *present)); + // If reflection succeeded, filter entities; otherwise keep all + if !reflection.is_empty() { + match Self::parse_reflection(&reflection) { + Ok(verified) => { + entities.retain(|e| verified.iter().any(|(name, present)| name == &e.name && *present)); + } + Err(e) => { + tracing::warn!("Reflection parse failed: {}, keeping all entities", e); + } + } + } else { + tracing::info!("Reflection skipped, keeping {} unverified entities", entities.len()); + } // Adjust confidence for reflected entities (slight penalty for needing verification) for entity in &mut entities {