From b8eb7efa6faf11f582a9091a7629c21a28aed12f Mon Sep 17 00:00:00 2001 From: rock Date: Thu, 10 Sep 2026 08:53:25 +0900 Subject: [PATCH] fix: Authentik JWT scope + ornith:35b reasoning field + observability logs Auth: - Add scope=openid roles to token request (required for llm:inference) - Derive TOKEN_URL from ISSUER or use TOKEN_URL env var - Support both AUTHENTIK_* and memory-agent-oidc secret key names ornith:35b support: - Handle reasoning field (content empty, JSON in reasoning) - Increase max_tokens to 12000 (reasoning models need headroom) - Fix trailing characters in fact extraction JSON parsing - Timeout increased to 120s for fact extraction Observability: - target="observability" structured logs for all LLM calls - event=llm_entity_call: model, endpoint, tokens, has_reasoning - event=llm_fact_call: model, endpoint, tokens, duration_ms - event=authentik_jwt_init: issuer, client_id - event=fact_jwt_fallback: error detail on JWT failure Column alignment: - INSERT uses source_id/target_id matching BFS query schema E2E tested with ornith:35b via api.riotpiao.com: - 6 entities, 4 edges with temporal facts - All observability logs present --- crates/mem-cli/src/ingest_worker.rs | 2 +- crates/mem-ingest/src/authentik_jwt.rs | 32 +++++++- crates/mem-ingest/src/entity_extractor.rs | 30 ++++--- crates/mem-ingest/src/fact_extractor.rs | 96 ++++++++++++++++++----- 4 files changed, 127 insertions(+), 33 deletions(-) diff --git a/crates/mem-cli/src/ingest_worker.rs b/crates/mem-cli/src/ingest_worker.rs index 9480af6..bc4a5bb 100644 --- a/crates/mem-cli/src/ingest_worker.rs +++ b/crates/mem-cli/src/ingest_worker.rs @@ -207,7 +207,7 @@ async fn save_entity_to_db(pool: &PgPool, entity: &mem_core::entity::Entity) -> async fn save_edge_to_db(pool: &PgPool, edge: &mem_core::edge::Edge) -> Result<()> { // Try temporal schema first (id, project_id, source_entity_id, etc) let result = sqlx::query( - "INSERT INTO memory_edge (id, project_id, source_entity_id, target_entity_id, relation_type, fact, t_valid, t_invalid, t_created, confidence) + "INSERT INTO memory_edge (id, project_id, source_id, target_id, relation_type, fact, t_valid, t_invalid, t_created, confidence) VALUES ($1, $2, $3, $4, $5, $6, $7::TIMESTAMPTZ, $8::TIMESTAMPTZ, $9::TIMESTAMPTZ, $10) ON CONFLICT (id) DO NOTHING" ) diff --git a/crates/mem-ingest/src/authentik_jwt.rs b/crates/mem-ingest/src/authentik_jwt.rs index f8b70bf..54b3e81 100644 --- a/crates/mem-ingest/src/authentik_jwt.rs +++ b/crates/mem-ingest/src/authentik_jwt.rs @@ -52,13 +52,24 @@ impl AuthentikJwtIssuer { /// From environment: AUTHENTIK_ISSUER, AUTHENTIK_CLIENT_ID, AUTHENTIK_CLIENT_SECRET pub fn from_env() -> Result { + // 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 diff --git a/crates/mem-ingest/src/entity_extractor.rs b/crates/mem-ingest/src/entity_extractor.rs index 271dcb9..5ad6e4e 100644 --- a/crates/mem-ingest/src/entity_extractor.rs +++ b/crates/mem-ingest/src/entity_extractor.rs @@ -156,7 +156,7 @@ impl LlmEntityExtractor { {"role": "user", "content": prompt} ], "temperature": 0.3, - "max_tokens": 1500 + "max_tokens": 12000 }); let response = client @@ -179,16 +179,28 @@ impl LlmEntityExtractor { } let data: serde_json::Value = response.json().await?; - let raw_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(); - // Strip ... tags from reasoning models - let content = Self::clean_llm_response(&raw_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); - tracing::debug!("LLM raw response length={}, cleaned length={}", raw_content.len(), content.len()); - tracing::debug!("LLM cleaned content: {}", content); + 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) } diff --git a/crates/mem-ingest/src/fact_extractor.rs b/crates/mem-ingest/src/fact_extractor.rs index 87cea1e..242ddab 100644 --- a/crates/mem-ingest/src/fact_extractor.rs +++ b/crates/mem-ingest/src/fact_extractor.rs @@ -82,11 +82,16 @@ impl FactExtractor for SimpleFactExtractor { /// Extracts relationships between entity pairs using LLM pub struct LlmFactExtractor { model_name: String, + jwt_issuer: Option>>, } impl LlmFactExtractor { pub fn new(model_name: &str) -> Self { - Self { model_name: model_name.to_string() } + 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 @@ -114,12 +119,27 @@ impl LlmFactExtractor { async fn call_llm(&self, prompt: &str) -> Result { let endpoint = std::env::var("LLM_ENDPOINT") - .unwrap_or_else(|_| "http://localhost:8081/v1/chat/completions".to_string()); + .unwrap_or_else(|_| "http://localhost:11434/v1/chat/completions".to_string()); - let api_key = std::env::var("LLM_API_KEY") - .or_else(|_| std::env::var("MEM_API_KEY")) - .unwrap_or_else(|_| "default-key".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, @@ -127,34 +147,50 @@ impl LlmFactExtractor { {"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": 1500, + "max_tokens": 12000, "temperature": 0.1 }); let response = client .post(&endpoint) - .header("Authorization", format!("Bearer {}", api_key)) + .header("Authorization", &auth_header) .header("Content-Type", "application/json") .json(&payload) - .timeout(std::time::Duration::from_secs(90)) + .timeout(std::time::Duration::from_secs(120)) .send() .await?; - if !response.status().is_success() { - let status = response.status(); + let status = response.status(); + if !status.is_success() { let body = response.text().await.unwrap_or_default(); - tracing::warn!("Fact extraction LLM error: {} - {}", status, body); + 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?; - let raw = data["choices"][0]["message"]["content"] - .as_str() - .unwrap_or("{}") - .to_string(); - let cleaned = Self::clean_llm_response(&raw); - tracing::debug!("Fact LLM response: raw_len={}, cleaned_len={}", raw.len(), cleaned.len()); + // 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) } } @@ -230,7 +266,29 @@ Respond in JSON: fact: String, } - match serde_json::from_str::(&response) { + // Try parsing, if trailing chars error try trimming to valid JSON + let parsed = match serde_json::from_str::(&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::(&response[..end]) + } else { + Err(e) + } + } + Err(e) => Err(e), + }; + match parsed { Ok(parsed) => { let facts: Vec = parsed.facts .into_iter() @@ -261,7 +319,7 @@ Respond in JSON: Ok(facts) } Err(e) => { - tracing::warn!("Fact extraction JSON parse failed: {}, response: {}", e, &response[..response.len().min(200)]); + tracing::warn!("Fact extraction JSON parse failed: {}", e); Ok(vec![]) } }