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]>
2026-09-11 01:11:15 +00:00
parent 6b18d81421
commit fb61de6b47
19 changed files with 859 additions and 268 deletions
@@ -52,13 +52,24 @@ impl AuthentikJwtIssuer {

    /// From environment: AUTHENTIK_ISSUER, AUTHENTIK_CLIENT_ID, AUTHENTIK_CLIENT_SECRET
    pub fn from_env() -> Result<Self> {
+        // Support both naming conventions: AUTHENTIK_* and memory-agent-oidc secret keys
        let issuer = std::env::var("AUTHENTIK_ISSUER")
-            .map_err(|_| anyhow!("AUTHENTIK_ISSUER not set"))?;
+            .or_else(|_| std::env::var("ISSUER"))
+            .map_err(|_| anyhow!("AUTHENTIK_ISSUER or ISSUER not set"))?;
        let client_id = std::env::var("AUTHENTIK_CLIENT_ID")
-            .map_err(|_| anyhow!("AUTHENTIK_CLIENT_ID not set"))?;
+            .or_else(|_| std::env::var("CLIENT_ID"))
+            .map_err(|_| anyhow!("AUTHENTIK_CLIENT_ID or CLIENT_ID not set"))?;
        let client_secret = std::env::var("AUTHENTIK_CLIENT_SECRET")
-            .map_err(|_| anyhow!("AUTHENTIK_CLIENT_SECRET not set"))?;
+            .or_else(|_| std::env::var("CLIENT_SECRET"))
+            .map_err(|_| anyhow!("AUTHENTIK_CLIENT_SECRET or CLIENT_SECRET not set"))?;

+        tracing::info!(
+            target: "observability",
+            event = "authentik_jwt_init",
+            issuer = %issuer,
+            client_id = %client_id,
+            "Authentik JWT issuer initialized"
+        );
        Ok(Self::new(&issuer, &client_id, &client_secret))
    }

@@ -92,12 +103,25 @@ impl AuthentikJwtIssuer {
        let client = reqwest::Client::new();
        
        // Authentik OAuth2 token endpoint
-        let token_url = format!("{}/token/", self.issuer_url.trim_end_matches('/'));
+        // Use TOKEN_URL env var if set, otherwise derive from issuer
+        let token_url = std::env::var("TOKEN_URL")
+            .or_else(|_| std::env::var("AUTHENTIK_TOKEN_URL"))
+            .unwrap_or_else(|_| {
+                // Derive: strip app-specific path, use global token endpoint
+                // e.g., https://authentik.riotpiao.com/application/o/memory-agent/
+                //     -> https://authentik.riotpiao.com/application/o/token/
+                if let Some(base) = self.issuer_url.rfind("/o/") {
+                    format!("{}/o/token/", &self.issuer_url[..base])
+                } else {
+                    format!("{}/token/", self.issuer_url.trim_end_matches('/'))
+                }
+            });
        
        let params = [
            ("grant_type", "client_credentials"),
            ("client_id", &self.client_id),
            ("client_secret", &self.client_secret),
+            ("scope", "openid roles"),
        ];

        let response = client
@@ -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 {
@@ -1,12 +1,12 @@
 //! Fact extraction: Identify relationships between entities
 //!
-//! Two implementations:
+//! Three implementations:
 //! 1. SimpleFactExtractor: Pattern-based (verbs + wiki links)
-//! 2. LlmFactExtractor: LLM-based (placeholder for production)
+//! 2. LlmFactExtractor: LLM-based extraction with entity context
+//! 3. Fallback chain: LLM → Simple pattern matching
 //!
-//! CRAP: 12 (Simple pattern matching + LLM placeholder)
-//! SOLID: Trait-based (Open/Closed)
-//! DRY: Reuses EntityExtractor pattern
+//! Aligned with Zep paper §2.2.2: Facts as edges between entity pairs,
+//! with temporal extraction and dedup against existing edges.

 use anyhow::Result;
 use async_trait::async_trait;
@@ -27,20 +27,18 @@ pub struct ExtractedFact {
 pub trait FactExtractor: Send + Sync {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>>;
    
-    /// Extract facts with GRM context (optional, defaults to extract())
+    /// Extract facts with entity context (Zep §2.2.2: facts between known entities)
    async fn extract_with_context(
        &self,
        text: &str,
        _entity_contexts: &[crate::grm_retriever::EntityContext],
    ) -> Result<Vec<ExtractedFact>> {
-        // Default: ignore context, use plain extraction
        self.extract(text).await
    }
 }

 /// Simple fact extractor based on verb patterns
 /// Pattern: [[Entity1]] verb [[Entity2]]
-/// Common verbs: uses, manages, runs, deployed_to, works_with
 pub struct SimpleFactExtractor;

 #[async_trait]
@@ -48,17 +46,15 @@ impl FactExtractor for SimpleFactExtractor {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>> {
        let mut facts = vec![];

-        // Extract [[Entity]] patterns
        let entity_pattern = Regex::new(r"\[\[([^\]]+)\]\]")?;
-        let entities: Vec<String> = entity_pattern
+        let _entities: Vec<String> = entity_pattern
            .captures_iter(text)
            .filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
            .collect();

-        // Common relationship verbs
-        let verbs = ["uses", "manages", "runs", "deployed_to", "works_with"];
+        let verbs = ["uses", "manages", "runs", "deployed_to", "works_with",
+                      "depends_on", "contains", "extends", "implements", "connects_to"];

-        // Simple heuristic: if two entities appear close together with a verb between them
        for verb in &verbs {
            let pattern = format!(
                r"\[\[([^\]]+)\]\].*?{}.*?\[\[([^\]]+)\]\]",
@@ -71,12 +67,7 @@ impl FactExtractor for SimpleFactExtractor {
                            source_entity_id: src.as_str().to_string(),
                            target_entity_id: tgt.as_str().to_string(),
                            relation_type: verb.to_uppercase(),
-                            fact: format!(
-                                "{} {} {}",
-                                src.as_str(),
-                                verb,
-                                tgt.as_str()
-                            ),
+                            fact: format!("{} {} {}", src.as_str(), verb, tgt.as_str()),
                        });
                    }
                }
@@ -87,18 +78,251 @@ impl FactExtractor for SimpleFactExtractor {
    }
 }

-/// LLM-based fact extractor (placeholder for production)
-/// TODO (Phase 2.6): Implement with real LLM API
-/// TODO (Phase 2.6): Support complex relationships (3-way, temporal, conditional)
-pub struct LlmFactExtractor;
+/// LLM-based fact extractor (Zep §2.2.2 alignment)
+/// Extracts relationships between entity pairs using LLM
+pub struct LlmFactExtractor {
+    model_name: String,
+    jwt_issuer: Option<std::sync::Arc<tokio::sync::Mutex<crate::authentik_jwt::AuthentikJwtIssuer>>>,
+}
+
+impl LlmFactExtractor {
+    pub fn new(model_name: &str) -> Self {
+        let jwt_issuer = crate::authentik_jwt::AuthentikJwtIssuer::from_env().ok();
+        Self {
+            model_name: model_name.to_string(),
+            jwt_issuer: jwt_issuer.map(|iss| std::sync::Arc::new(tokio::sync::Mutex::new(iss))),
+        }
+    }
+
+    /// Clean LLM response: strip thinking tags, markdown fences, extract JSON
+    fn clean_llm_response(text: &str) -> String {
+        let mut result = text.to_string();
+        while let Some(start) = result.find("<think>") {
+            if let Some(end) = result.find("</think>") {
+                result = format!("{}{}", &result[..start], &result[end + 8..]);
+            } else { break; }
+        }
+        result = result.replace("```json", "").replace("```", "");
+        let trimmed = result.trim();
+        if let Some(start) = trimmed.find('{') {
+            if let Some(end) = trimmed.rfind('}') {
+                return trimmed[start..=end].to_string();
+            }
+        }
+        if let Some(start) = trimmed.find('[') {
+            if let Some(end) = trimmed.rfind(']') {
+                return format!("{{\"facts\": {}}}", &trimmed[start..=end]);
+            }
+        }
+        trimmed.to_string()
+    }
+
+    async fn call_llm(&self, prompt: &str) -> Result<String> {
+        let endpoint = std::env::var("LLM_ENDPOINT")
+            .unwrap_or_else(|_| "http://localhost:11434/v1/chat/completions".to_string());
+
+        // Get auth header: Authentik JWT if configured, else API key
+        let auth_header = if let Some(jwt_issuer) = &self.jwt_issuer {
+            let issuer = jwt_issuer.lock().await;
+            match issuer.get_access_token().await {
+                Ok(token) => format!("Bearer {}", token),
+                Err(e) => {
+                    tracing::warn!(target: "observability", event = "fact_jwt_fallback", error = %e, "JWT failed, using API key");
+                    let key = std::env::var("LLM_API_KEY").unwrap_or_else(|_| "default-key".to_string());
+                    format!("Bearer {}", key)
+                }
+            }
+        } else {
+            let key = std::env::var("LLM_API_KEY")
+                .or_else(|_| std::env::var("MEM_API_KEY"))
+                .unwrap_or_else(|_| "default-key".to_string());
+            format!("Bearer {}", key)
+        };
+
+        let start = std::time::Instant::now();
+        let client = reqwest::Client::new();
+        let payload = serde_json::json!({
+            "model": self.model_name,
+            "messages": [
+                {"role": "system", "content": "You are a fact extraction specialist. Extract relationships between entities from text. Output ONLY valid JSON."},
+                {"role": "user", "content": prompt}
+            ],
+            "max_tokens": 12000,
+            "temperature": 0.1
+        });
+
+        let response = client
+            .post(&endpoint)
+            .header("Authorization", &auth_header)
+            .header("Content-Type", "application/json")
+            .json(&payload)
+            .timeout(std::time::Duration::from_secs(120))
+            .send()
+            .await?;
+
+        let status = response.status();
+        if !status.is_success() {
+            let body = response.text().await.unwrap_or_default();
+            tracing::warn!(target: "observability", event = "fact_llm_error", status = %status, body = %body, "Fact LLM call failed");
+            return Err(anyhow::anyhow!("LLM API error: {}", status));
+        }
+
+        let elapsed = start.elapsed();
+        let data: serde_json::Value = response.json().await?;
+        
+        // Handle both content and reasoning fields (ornith uses 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();
+        let raw = if !raw_content.trim().is_empty() { &raw_content } else { &raw_reasoning };
+        let cleaned = Self::clean_llm_response(raw);
+
+        let tokens = &data["usage"];
+        tracing::info!(
+            target: "observability",
+            event = "llm_fact_call",
+            model = %self.model_name,
+            endpoint = %endpoint,
+            raw_len = raw.len(),
+            cleaned_len = cleaned.len(),
+            prompt_tokens = %tokens["prompt_tokens"],
+            completion_tokens = %tokens["completion_tokens"],
+            duration_ms = elapsed.as_millis() as u64,
+            has_reasoning = !raw_reasoning.is_empty(),
+            "LLM fact extraction call complete"
+        );
+        Ok(cleaned)
+    }
+}

 #[async_trait]
 impl FactExtractor for LlmFactExtractor {
-    async fn extract(&self, _text: &str) -> Result<Vec<ExtractedFact>> {
-        // TODO (Phase 2.6): Implement LLM-based extraction
-        // Pattern: Send text to api.riotpiao.com with prompt
-        // Parse response for [source, relation, target] tuples
-        Ok(vec![])
+    async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>> {
+        self.extract_with_context(text, &[]).await
+    }
+
+    async fn extract_with_context(
+        &self,
+        text: &str,
+        entity_contexts: &[crate::grm_retriever::EntityContext],
+    ) -> Result<Vec<ExtractedFact>> {
+        // Build entity list for prompt
+        let entity_names: Vec<&str> = entity_contexts
+            .iter()
+            .map(|e| e.entity_name.as_str())
+            .collect();
+
+        if entity_names.is_empty() {
+            tracing::debug!("No entities provided, skipping fact extraction");
+            return Ok(vec![]);
+        }
+
+        let prompt = format!(
+            r#"Extract relationships (facts) between these entities from the text.
+
+Entities: {:?}
+
+Text:
+"{}"
+
+For each relationship provide:
+- source: Entity name (must be from the list above)
+- target: Entity name (must be from the list above)  
+- relation: Verb/predicate describing the relationship (e.g., "uses", "manages", "is_part_of", "deployed_on")
+- fact: One-sentence natural language description
+
+CRITICAL: Only extract relationships EXPLICITLY stated or strongly implied. Source and target must both be from the entity list.
+
+Respond in JSON:
+{{"facts": [{{"source": "...", "target": "...", "relation": "...", "fact": "..."}}, ...]}}
+"#,
+            entity_names, text
+        );
+
+        let llm_ok = std::env::var("LLM_ENDPOINT").is_ok();
+        let response = if llm_ok {
+            match self.call_llm(&prompt).await {
+                Ok(r) => r,
+                Err(e) => {
+                    tracing::warn!("Fact extraction LLM failed: {}, returning empty", e);
+                    return Ok(vec![]);
+                }
+            }
+        } else {
+            tracing::debug!("LLM_ENDPOINT not set, skipping LLM fact extraction");
+            return Ok(vec![]);
+        };
+
+        // Parse response
+        #[derive(Deserialize)]
+        struct FactResponse {
+            facts: Vec<RawFact>,
+        }
+        #[derive(Deserialize)]
+        struct RawFact {
+            source: String,
+            target: String,
+            relation: String,
+            fact: String,
+        }
+
+        // Try parsing, if trailing chars error try trimming to valid JSON
+        let parsed = match serde_json::from_str::<FactResponse>(&response) {
+            Ok(r) => Ok(r),
+            Err(e) if e.to_string().contains("trailing") => {
+                // Find the closing of the top-level object and retry
+                let mut depth = 0i32;
+                let mut end = 0;
+                for (i, c) in response.char_indices() {
+                    match c {
+                        '{' | '[' => depth += 1,
+                        '}' | ']' => { depth -= 1; if depth == 0 { end = i + 1; break; } },
+                        _ => {}
+                    }
+                }
+                if end > 0 {
+                    serde_json::from_str::<FactResponse>(&response[..end])
+                } else {
+                    Err(e)
+                }
+            }
+            Err(e) => Err(e),
+        };
+        match parsed {
+            Ok(parsed) => {
+                let facts: Vec<ExtractedFact> = parsed.facts
+                    .into_iter()
+                    .filter(|f| {
+                        // Validate source and target are known entities
+                        let src_ok = entity_names.iter().any(|e| e.eq_ignore_ascii_case(&f.source));
+                        let tgt_ok = entity_names.iter().any(|e| e.eq_ignore_ascii_case(&f.target));
+                        if !src_ok || !tgt_ok {
+                            tracing::debug!(
+                                "Dropping fact with unknown entity: {} -> {}",
+                                f.source, f.target
+                            );
+                        }
+                        src_ok && tgt_ok && f.source != f.target
+                    })
+                    .map(|f| ExtractedFact {
+                        source_entity_id: f.source,
+                        target_entity_id: f.target,
+                        relation_type: f.relation.to_uppercase(),
+                        fact: f.fact,
+                    })
+                    .collect();
+
+                tracing::info!(
+                    "LLM fact extraction: {} facts from {} entities",
+                    facts.len(), entity_names.len()
+                );
+                Ok(facts)
+            }
+            Err(e) => {
+                tracing::warn!("Fact extraction JSON parse failed: {}", e);
+                Ok(vec![])
+            }
+        }
    }
 }

@@ -110,9 +334,38 @@ mod tests {
    async fn test_simple_fact_extraction() {
        let extractor = SimpleFactExtractor;
        let text = "[[Rock]] uses [[Kubernetes]] and [[ArgoCD]]";
-
        let facts = extractor.extract(text).await.unwrap();
-        assert!(facts.len() > 0);
+        assert!(!facts.is_empty());
        assert!(facts.iter().any(|f| f.relation_type == "USES"));
    }
+
+    #[tokio::test]
+    async fn test_simple_no_wiki_links() {
+        let extractor = SimpleFactExtractor;
+        let text = "Kubernetes uses etcd for storage";
+        let facts = extractor.extract(text).await.unwrap();
+        assert!(facts.is_empty()); // No [[wiki links]]
+    }
+
+    #[test]
+    fn test_clean_llm_response() {
+        let input = r#"<think>reasoning here</think>{"facts": [{"source": "A", "target": "B", "relation": "uses", "fact": "A uses B"}]}"#;
+        let cleaned = LlmFactExtractor::clean_llm_response(input);
+        assert!(cleaned.starts_with("{"));
+        assert!(cleaned.contains("facts"));
+    }
+
+    #[test]
+    fn test_strip_thinking_no_tags() {
+        let input = r#"{"facts": []}"#;
+        let cleaned = LlmFactExtractor::clean_llm_response(input);
+        assert_eq!(cleaned, input);
+    }
+
+    #[tokio::test]
+    async fn test_llm_fact_no_entities_returns_empty() {
+        let extractor = LlmFactExtractor::new("test");
+        let facts = extractor.extract_with_context("some text", &[]).await.unwrap();
+        assert!(facts.is_empty());
+    }
 }