feat: LLM entity + fact extraction pipeline (Zep paper alignment) (#48)
## 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:
@@ -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,35 @@ impl LlmEntityExtractor {
|
||||
|
||||
/// Parse extraction response JSON
|
||||
/// Format: { "entities": [{ "name": "...", "type": "...", "summary": "..." }, ...] }
|
||||
/// Clean LLM response: strip thinking tags, markdown fences, extract JSON
|
||||
fn clean_llm_response(text: &str) -> String {
|
||||
let mut result = text.to_string();
|
||||
// Remove <think>...</think> blocks
|
||||
while let Some(start) = result.find("<think>") {
|
||||
if let Some(end) = result.find("</think>") {
|
||||
result = format!("{}{}", &result[..start], &result[end + 8..]);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Remove markdown code fences
|
||||
result = result.replace("```json", "").replace("```", "");
|
||||
// Find JSON object
|
||||
let trimmed = result.trim();
|
||||
if let Some(start) = trimmed.find('{') {
|
||||
if let Some(end) = trimmed.rfind('}') {
|
||||
return trimmed[start..=end].to_string();
|
||||
}
|
||||
}
|
||||
// Maybe it's a JSON array — wrap in object
|
||||
if let Some(start) = trimmed.find('[') {
|
||||
if let Some(end) = trimmed.rfind(']') {
|
||||
return format!("{{\"entities\": {}}}", &trimmed[start..=end]);
|
||||
}
|
||||
}
|
||||
trimmed.to_string()
|
||||
}
|
||||
|
||||
fn parse_extraction(response: &str) -> Result<Vec<ExtractedEntity>> {
|
||||
#[derive(Deserialize)]
|
||||
struct Response {
|
||||
@@ -123,7 +156,7 @@ impl LlmEntityExtractor {
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
"temperature": 0.3,
|
||||
"max_tokens": 500
|
||||
"max_tokens": 12000
|
||||
});
|
||||
|
||||
let response = client
|
||||
@@ -131,7 +164,7 @@ impl LlmEntityExtractor {
|
||||
.header("Authorization", auth_header)
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&payload)
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.timeout(std::time::Duration::from_secs(90))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
@@ -146,12 +179,28 @@ impl LlmEntityExtractor {
|
||||
}
|
||||
|
||||
let data: serde_json::Value = response.json().await?;
|
||||
let content = data["choices"][0]["message"]["content"]
|
||||
.as_str()
|
||||
.unwrap_or("{}")
|
||||
.to_string();
|
||||
// Extract content — some models put JSON in "content", others in "reasoning"
|
||||
let msg = &data["choices"][0]["message"];
|
||||
let raw_content = msg["content"].as_str().unwrap_or("").to_string();
|
||||
let raw_reasoning = msg["reasoning"].as_str().unwrap_or("").to_string();
|
||||
|
||||
tracing::debug!("LLM response (via Authentik JWT): {}", content);
|
||||
// Use content if non-empty, otherwise try reasoning field
|
||||
let raw = if !raw_content.trim().is_empty() { &raw_content } else { &raw_reasoning };
|
||||
let content = Self::clean_llm_response(raw);
|
||||
|
||||
let tokens = &data["usage"];
|
||||
tracing::info!(
|
||||
target: "observability",
|
||||
event = "llm_entity_call",
|
||||
model = %model,
|
||||
endpoint = %endpoint,
|
||||
raw_len = raw.len(),
|
||||
cleaned_len = content.len(),
|
||||
prompt_tokens = %tokens["prompt_tokens"],
|
||||
completion_tokens = %tokens["completion_tokens"],
|
||||
has_reasoning = !raw_reasoning.is_empty(),
|
||||
"LLM entity extraction call complete"
|
||||
);
|
||||
Ok(content)
|
||||
}
|
||||
|
||||
@@ -233,14 +282,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