crates/mem-ingest/src/entity_extractor.rs

@@ -44,6 +44,10 @@ impl ExtractedEntity {
 #[async_trait]
 pub trait EntityExtractor: Send + Sync {
     async fn extract(&self, text: &str) -> Result<Vec<ExtractedEntity>>;
+    async fn extract_with_auth(&self, text: &str, x_forward_user: Option<&str>) -> Result<Vec<ExtractedEntity>> {
+        // Default: ignore auth header, use regular extract
+        self.extract(text).await
+    }
 }
 
 /// LLM-based extractor with reflection verification (stage 1 + 2)
@@ -121,29 +125,42 @@ impl LlmEntityExtractor {
         Ok(parsed.verified.into_iter().map(|v| (v.name, v.present)).collect())
     }
 
-    /// Call LLM via api.riotpiao.com using Authentik JWT
-    /// Token is fetched from Authentik service account and cached
-    async fn call_llm_endpoint(&self, prompt: &str) -> Result<String> {
+    /// Call LLM via api.riotpiao.com using X-Forward-User auth/exchange
+    /// Supports: Authentik JWT, X-Forward-User header, or API key fallback
+    async fn call_llm_endpoint(&self, prompt: &str, x_forward_user: Option<&str>) -> Result<String> {
         let endpoint = std::env::var("LLM_ENDPOINT")
             .unwrap_or_else(|_| "http://api-internal.riotpiao.com:8000/v1/chat/completions".to_string());
         let model = std::env::var("LLM_MODEL")
             .unwrap_or_else(|_| "qwen:7b".to_string());
         
-        // Get JWT token from Authentik
-        let auth_header = if let Some(jwt_issuer) = &self.jwt_issuer {
+        // Get auth header: prefer X-Forward-User, fallback to Authentik JWT, then API key
+        let auth_header = if let Some(user) = x_forward_user {
+            // Use X-Forward-User directly (API Gateway pattern)
+            tracing::info!("Using X-Forward-User for LLM auth: {}", user);
+            format!("X-Forward-User: {}", user)
+        } else 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),
+                Ok(token) => {
+                    tracing::info!("Using Authentik JWT for LLM auth");
+                    format!("Bearer {}", token)
+                },
                 Err(e) => {
                     tracing::warn!("Failed to get Authentik JWT: {}", e);
-                    return Err(e);
+                    // Fallback to env var
+                    let api_key = std::env::var("LLM_API_KEY")
+                        .or_else(|_| std::env::var("MEM_API_KEY"))
+                        .unwrap_or_else(|_| "test-key".to_string());
+                    tracing::info!("Falling back to LLM_API_KEY");
+                    format!("Bearer {}", api_key)
                 }
             }
         } else {
             // Fallback to env var if Authentik not configured
             let api_key = std::env::var("LLM_API_KEY")
                 .or_else(|_| std::env::var("MEM_API_KEY"))
-                .unwrap_or_else(|_| "default-key".to_string());
+                .unwrap_or_else(|_| "test-key".to_string());
+            tracing::info!("Using LLM_API_KEY for LLM auth");
             format!("Bearer {}", api_key)
         };
         
@@ -160,23 +177,33 @@ impl LlmEntityExtractor {
             "max_tokens": 12000
         });
         
-        let response = client
+        let mut request = client
             .post(&endpoint)
-            .header("Authorization", auth_header)
-            .header("Content-Type", "application/json")
+            .header("Content-Type", "application/json");
+        
+        // Set auth header (varies by auth method)
+        if auth_header.starts_with("X-Forward-User") {
+            request = request.header("X-Forward-User", auth_header.split(": ").nth(1).unwrap_or("unknown"));
+        } else {
+            request = request.header("Authorization", auth_header);
+        }
+        
+        let response = request
             .json(&payload)
             .timeout(std::time::Duration::from_secs(90))
             .send()
             .await?;
         
-        if !response.status().is_success() {
-            tracing::warn!(
+        let status = response.status();
+        if !status.is_success() {
+            let error_text = response.text().await.unwrap_or_default();
+            tracing::error!(
                 "LLM API error: {} - {}",
-                response.status(),
-                response.text().await.unwrap_or_default()
+                status,
+                error_text
             );
-            // Fallback to mock response on error
-            return Ok(r#"{"entities": []}"#.to_string());
+            // Return error instead of silently returning empty array
+            return Err(anyhow::anyhow!("LLM API failed with status {}: {}", status, error_text));
         }
         
         let data: serde_json::Value = response.json().await?;
@@ -258,7 +285,10 @@ Respond in JSON:
 
         // Try real LLM first, fallback to mock if not configured
         let extraction_response = if std::env::var("LLM_ENDPOINT").is_ok() {
-            self.call_llm_endpoint(&prompt).await.unwrap_or_else(|_| self.simulate_llm(&prompt).unwrap_or_default())
+            self.call_llm_endpoint(&prompt, None).await.unwrap_or_else(|e| {
+                tracing::error!("LLM entity extraction failed: {}, using mock", e);
+                self.simulate_llm(&prompt).unwrap_or_default()
+            })
         } else {
             self.simulate_llm(&prompt)?
         };
@@ -283,7 +313,7 @@ Respond in JSON:
             );
 
             let reflection = if std::env::var("LLM_ENDPOINT").is_ok() {
-                self.call_llm_endpoint(&reflection_prompt).await.unwrap_or_else(|e| {
+                self.call_llm_endpoint(&reflection_prompt, None).await.unwrap_or_else(|e| {
                     tracing::warn!("Reflection LLM call failed: {}, skipping verification", e);
                     String::new()
                 })
@@ -313,6 +343,85 @@ Respond in JSON:
 
         Ok(entities)
     }
+
+    /// Extract with X-Forward-User auth header (API Gateway pattern)
+    async fn extract_with_auth(&self, text: &str, x_forward_user: Option<&str>) -> Result<Vec<ExtractedEntity>> {
+        let mut entities = vec![];
+        
+        // Extract speaker if available
+        use crate::speaker_extractor::{HeuristicSpeakerExtractor, SpeakerConfig};
+        if let Ok(speaker_extractor) = HeuristicSpeakerExtractor::new(SpeakerConfig::default()) {
+            if let Ok(Some(speaker)) = speaker_extractor.extract_speaker(text).await {
+                entities.push(ExtractedEntity {
+                    name: speaker.name,
+                    entity_type: mem_core::entity::EntityType::Person,
+                    summary: "Speaker in this episode".to_string(),
+                    confidence: speaker.confidence,
+                });
+            }
+        }
+
+        // Extract entities with auth header
+        let prompt = format!(
+            r#"Extract named entities from this text.
+
+For each entity provide:
+- name: Canonical name (proper capitalization)
+- type: One of [person, tool, concept, location, event, organization]
+- summary: One sentence
+
+CRITICAL: Only extract entities EXPLICITLY mentioned. No inference.
+
+Text:
+"{}"
+
+Respond in JSON:
+{{"entities": [{{"name": "...", "type": "...", "summary": "..."}}, ...]}}
+"#,
+            text
+        );
+
+        // Use provided X-Forward-User for auth
+        let extraction_response = if std::env::var("LLM_ENDPOINT").is_ok() {
+            self.call_llm_endpoint(&prompt, x_forward_user).await.unwrap_or_else(|e| {
+                tracing::error!("LLM entity extraction with auth failed: {}", e);
+                self.simulate_llm(&prompt).unwrap_or_default()
+            })
+        } else {
+            self.simulate_llm(&prompt)?
+        };
+        
+        let extracted = Self::parse_extraction(&extraction_response)?;
+        entities.extend(extracted);
+
+        // Optional: reflection verification with auth
+        if self.enable_reflection && std::env::var("LLM_ENDPOINT").is_ok() {
+            let reflection_prompt = format!(
+                r#"Verify these entities are explicitly in the text:
+
+Text:
+"{}"
+
+Entities:
+{:?}
+
+Respond in JSON:
+{{"verified": [{{"name": "...", "present": true/false}}, ...]}}
+"#,
+                text, entities
+            );
+
+            if let Ok(reflection) = self.call_llm_endpoint(&reflection_prompt, x_forward_user).await {
+                if !reflection.is_empty() {
+                    if let Ok(verified) = Self::parse_reflection(&reflection) {
+                        entities.retain(|e| verified.iter().any(|(name, present)| name == &e.name && *present));
+                    }
+                }
+            }
+        }
+
+        Ok(entities)
+    }
 }
 
 /// Fallback extractor: Use wiki_links if LLM fails (stage 3)