fix: enable LLM entity extraction + handle reasoning model output
Root causes of zero entity extraction: 1. IngestWorker used WikiLinkFallbackExtractor (wiki links only) Fix: Use LlmEntityExtractor when LLM_ENDPOINT is set 2. ExtractedEntity.entity_type vs LLM returning "type" Fix: serde alias "type" -> entity_type, default confidence 3. Reasoning models output <think>...</think> before JSON Fix: strip_thinking_tags() extracts JSON from response 4. Reflection verification crashes pipeline on parse failure Fix: graceful fallback, keep all entities if reflection fails Tested with reasoning-predictor (qwen2.5:3b) via port-forward.
This commit is contained in:
@@ -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<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);
|
||||
let contradiction_detector = Arc::new(ContradictionHandler::default());
|
||||
|
||||
@@ -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 <think>...</think> tags from reasoning model output and extract JSON
|
||||
fn strip_thinking_tags(text: &str) -> String {
|
||||
let mut result = text.to_string();
|
||||
// Remove <think>...</think> blocks
|
||||
if let Some(start) = result.find("<think>") {
|
||||
if let Some(end) = result.find("</think>") {
|
||||
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<Vec<ExtractedEntity>> {
|
||||
#[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 <think>...</think> 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 {
|
||||
|
||||
Reference in New Issue
Block a user