fix: Authentik JWT scope + ornith:35b reasoning field + observability logs
CI / CI (pull_request) Successful in 11m57s
CI / CI (pull_request) Successful in 11m57s
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
This commit is contained in:
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 <think>...</think> 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)
|
||||
}
|
||||
|
||||
|
||||
@@ -82,11 +82,16 @@ impl FactExtractor for SimpleFactExtractor {
|
||||
/// 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 {
|
||||
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<String> {
|
||||
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::<FactResponse>(&response) {
|
||||
// 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()
|
||||
@@ -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![])
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user