klieo-memory-graph-rag 2.1.0

//! `LlmEntityExtractor` — recall-extension `EntityExtractor` that uses
//! an LLM to recognise entities in free-form text.
//!
//! Wired as the secondary leg of [`FallbackExtractor`](crate::FallbackExtractor)
//! with [`BuiltinExtractor`](crate::BuiltinExtractor) as the primary. The
//! LLM fires only on text the regex couldn't recognise.
//!
//! **Error contract:** LLM failures (transport error, model unavailable,
//! script exhausted in fakes) AND parse failures (non-JSON output, schema
//! mismatch) BOTH fall back silently to caller hints. Recall is
//! best-effort; never propagate provider failures to the ingest path
//! since graph indexing is non-authoritative (the vector store is the
//! source of truth in `GraphAwareLongTerm::remember`).

use async_trait::async_trait;
use klieo_core::error::MemoryError;
use klieo_core::llm::{ChatRequest, LlmClient, Message, ResponseFormat, Role};
use klieo_memory_graph::{EntityExtractor, EntityRef, EntityType};
use serde::Deserialize;
use std::collections::HashSet;
use std::sync::Arc;

const SYSTEM_PROMPT: &str = "You are an entity-extraction service. Given a text body, \
identify typed entities and return them as a JSON array. Each element MUST be an object with \
exactly two string fields: `type` (one of: Ticket, Policy, Agent, Member, Concept; or a free-form \
domain label for novel categories) and `name` (the entity identifier as it appears in the text, \
preserving case). Return ONLY the JSON array — no prose, no preamble. An empty array (`[]`) means \
no entities were found.";

/// LLM-driven `EntityExtractor`. Prompts a `LlmClient` for a JSON-array
/// response describing typed entities in the input text.
pub struct LlmEntityExtractor {
    client: Arc<dyn LlmClient>,
}

impl LlmEntityExtractor {
    /// Build from any `Arc<dyn LlmClient>` — provider-agnostic.
    pub fn new(client: Arc<dyn LlmClient>) -> Self {
        Self { client }
    }

    fn build_request(text: &str) -> ChatRequest {
        let messages = vec![
            Message {
                role: Role::System,
                content: SYSTEM_PROMPT.to_string(),
                tool_calls: Vec::new(),
                tool_call_id: None,
            },
            Message {
                role: Role::User,
                content: text.to_string(),
                tool_calls: Vec::new(),
                tool_call_id: None,
            },
        ];
        ChatRequest {
            response_format: ResponseFormat::Json {
                schema: schema_for_entity_array(),
            },
            ..ChatRequest::new(messages)
        }
    }
}

#[derive(Debug, Deserialize)]
struct WireEntity {
    #[serde(rename = "type")]
    entity_type: String,
    name: String,
}

fn schema_for_entity_array() -> serde_json::Value {
    serde_json::json!({
        "type": "array",
        "items": {
            "type": "object",
            "required": ["type", "name"],
            "properties": {
                "type": { "type": "string" },
                "name": { "type": "string" }
            }
        }
    })
}

fn parse_response(body: &str) -> Option<Vec<EntityRef>> {
    let parsed: Vec<WireEntity> = serde_json::from_str(body.trim()).ok()?;
    let mut out = Vec::with_capacity(parsed.len());
    for w in parsed {
        let entity_type = EntityType::from_label(&w.entity_type);
        // try_new rejects Other{hint} that collides with a reserved
        // label. Drop with a WARN so recall loss surfaces in ops logs
        // rather than silently corrupting the entity index — recall is
        // best-effort and never blocks ingest.
        match EntityRef::try_new(entity_type, w.name.clone()) {
            Ok(entity_ref) => out.push(entity_ref),
            Err(err) => {
                tracing::warn!(
                    extractor = "LlmEntityExtractor",
                    label = %w.entity_type,
                    name = %w.name,
                    error = %err,
                    "LLM-emitted entity rejected (reserved-label collision); dropping"
                );
            }
        }
    }
    Some(out)
}

#[async_trait]
impl EntityExtractor for LlmEntityExtractor {
    async fn extract(
        &self,
        text: &str,
        hints: &[EntityRef],
    ) -> Result<Vec<EntityRef>, MemoryError> {
        let request = Self::build_request(text);
        let response = match self.client.complete(request).await {
            Ok(r) => r,
            Err(e) => {
                tracing::warn!(
                    extractor = "LlmEntityExtractor",
                    error = %e,
                    "LLM call failed; falling back to caller hints"
                );
                return Ok(hints.to_vec());
            }
        };

        let llm_entities = match parse_response(&response.message.content) {
            Some(es) => es,
            None => {
                tracing::warn!(
                    extractor = "LlmEntityExtractor",
                    body = %response.message.content,
                    "LLM response was not a parseable JSON entity array; falling back to caller hints"
                );
                return Ok(hints.to_vec());
            }
        };

        // Merge hints + LLM hits, dedup by (EntityType::as_str(), name).
        let mut seen: HashSet<(String, String)> = HashSet::new();
        let mut out = Vec::with_capacity(hints.len() + llm_entities.len());
        for h in hints {
            let key = (h.entity_type.as_str().to_owned(), h.name.clone());
            if seen.insert(key) {
                out.push(h.clone());
            }
        }
        for h in llm_entities {
            let key = (h.entity_type.as_str().to_owned(), h.name.clone());
            if seen.insert(key) {
                out.push(h);
            }
        }
        Ok(out)
    }
}