polyc_agent/

extraction.rs

//! The provider-agnostic post-turn memory extractor
//! (docs/design/personas.md §5, issue #214).
//!
//! After a turn commits, a cheap model distills the transcript into durable
//! facts about the person — never in the turn's hot path. The extractor is
//! handed the persona's existing ACTIVE facts so contradictions come back as
//! invalidations (paired with the replacing fact) rather than duplicates;
//! the store then closes the old fact's validity interval — invalidate,
//! never delete.
//!
//! Like the participation classifier next door, this is deliberately
//! provider-agnostic and **tolerant on the way out**: a reply that isn't
//! the expected JSON extracts nothing (memory is best-effort enrichment; a
//! flaky cheap model must never break the pipeline).

use polyc_llm::{CompletionRequest, Content, LlmProvider, Message, Role, turn::collect_turn};
use serde::Deserialize;

use crate::participation::ParticipationMsg;

/// Most facts one turn may add. A chatty turn distills to a few durable
/// facts; dozens means the model is transcribing, not distilling.
pub const MAX_FACTS_PER_TURN: usize = 8;

/// Write-time confidence floor (`#796`).
///
/// Basis points, the same scale as [`CandidateFact::confidence_bps`]: a fact
/// the classifier itself is not reasonably sure of must never become a
/// durable, cross-conversation "fact" just because the model emitted
/// well-formed JSON. 6000 = 60%.
pub const MIN_CONFIDENCE_BPS: u32 = 6_000;

/// One fact the extractor proposes to remember.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CandidateFact {
    /// The fact, one self-contained sentence.
    pub text: String,
    /// Entities the fact mentions.
    pub entities: Vec<String>,
    /// Extractor confidence in basis points (0–10000).
    pub confidence_bps: u32,
    /// The id of the existing fact this one supersedes, when the extractor
    /// paired the add with a contradiction — what lets the store link the
    /// closed interval to its replacement (`superseded_by`).
    pub replaces: Option<String>,
}

/// One existing fact the extractor proposes to invalidate (contradicted by
/// this turn).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Invalidation {
    /// The id of the existing fact whose validity interval should close.
    pub fact_id: String,
    /// Why (one short phrase).
    pub reason: String,
}

/// What one extraction pass proposes.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ExtractedMemories {
    /// New facts to append.
    pub added: Vec<CandidateFact>,
    /// Existing facts this turn contradicted.
    pub invalidated: Vec<Invalidation>,
    /// Ids of existing facts this turn merely RESTATED — the same claim in
    /// different words, no new information. The semantic merge pass (#860) folds
    /// each into its existing entry as a corroboration instead of appending a
    /// near-duplicate (INV-P26), never crossing scope (INV-P24). Empty when the
    /// turn restated nothing.
    pub corroborated: Vec<String>,
}

/// An existing active fact, as shown to the extractor for contradiction
/// checks.
#[derive(Debug, Clone)]
pub struct ExistingFact {
    /// The fact's journal id (what an invalidation must reference).
    pub fact_id: String,
    /// The fact text.
    pub text: String,
}

/// The extractor's JSON reply shape (tolerantly deserialized).
#[derive(Debug, Default, Deserialize)]
struct WireReply {
    #[serde(default)]
    added: Vec<WireFact>,
    #[serde(default)]
    invalidated: Vec<WireInvalidation>,
    #[serde(default)]
    corroborated: Vec<String>,
}

#[derive(Debug, Deserialize)]
struct WireFact {
    #[serde(default)]
    text: String,
    #[serde(default)]
    entities: Vec<String>,
    /// 0–100; clamped and scaled to basis points.
    #[serde(default)]
    confidence: u32,
    /// The id of the existing fact this one replaces; empty when the add is
    /// not a replacement.
    #[serde(default)]
    replaces: String,
}

#[derive(Debug, Deserialize)]
struct WireInvalidation {
    #[serde(default)]
    fact_id: String,
    #[serde(default)]
    reason: String,
}

/// System prompt for the extraction pass.
const fn system_prompt() -> &'static str {
    "You distill a conversation turn into durable facts about the person speaking — things \
     worth remembering across future conversations (preferences, role, projects, standing \
     constraints). Ignore small talk, one-off logistics, and anything about the assistant \
     itself. Never emit a fact naming a home address, a phone number, a government id \
     (SSN, passport, driver's license), a financial account or card number, a password or \
     API/secret key, or a health/medical detail — omit the fact entirely rather than \
     write around it. Set confidence honestly (0-100): a fact you are not reasonably sure \
     of is worse than no fact, so lean low rather than guess. You are also given the \
     person's EXISTING facts with ids; when this turn contradicts one, list its id under \
     invalidated AND add the replacement fact under added with \"replaces\" set to that \
     same id, so the old fact links to its replacement. Omit \"replaces\" for a fact that \
     replaces nothing. When this turn merely RESTATES an existing fact — the same claim in \
     different words, with no new or changed information — do NOT add it: list that existing \
     fact's id under corroborated instead, so the known fact is reinforced rather than \
     duplicated.\n\
     Reply with ONLY this JSON, no prose:\n\
     {\"added\":[{\"text\":\"…\",\"entities\":[\"…\"],\"confidence\":0-100,\
     \"replaces\":\"existing fact id, or omit\"}],\
     \"invalidated\":[{\"fact_id\":\"…\",\"reason\":\"…\"}],\
     \"corroborated\":[\"existing fact id\"]}\n\
     All arrays may be empty. At most a few added facts per turn."
}

/// Render the extractor's user message: existing facts (with ids), then the
/// turn transcript.
fn render_input(transcript: &[ParticipationMsg], existing: &[ExistingFact]) -> String {
    use std::fmt::Write as _;
    let mut out = String::new();
    out.push_str("EXISTING FACTS:\n");
    if existing.is_empty() {
        out.push_str("(none)\n");
    }
    for fact in existing {
        let _ = writeln!(out, "- [{}] {}", fact.fact_id, fact.text);
    }
    out.push_str("\nTURN TRANSCRIPT:\n");
    for msg in transcript {
        let speaker = if msg.is_self {
            "assistant"
        } else {
            &msg.speaker
        };
        out.push_str(speaker);
        out.push_str(": ");
        out.push_str(&msg.text);
        out.push('\n');
    }
    out
}

/// Refused-category keywords for the post-parse PII heuristic (`#796`):
/// lowercased substrings that, anywhere in a candidate fact's text, mean the
/// fact names an identifier or category durable memory must never carry —
/// a home address, a health/medical detail, or a credential/identifier.
/// Defense in depth: the extractor prompt already instructs the model to
/// omit these, but a model that ignores the instruction gets no second
/// chance at the parser. Deliberately over-inclusive (a false positive costs
/// one skipped fact; a false negative persists PII).
const PII_REFUSAL_KEYWORDS: &[&str] = &[
    "ssn",
    "social security",
    "credit card",
    "card number",
    "cvv",
    "passport number",
    "driver's license",
    "password",
    "api key",
    "secret key",
    "private key",
    "home address",
    "lives at",
    "street address",
    "diagnosed with",
    "medical condition",
    "prescription",
    "medication",
    "mental health",
];

/// Whether `text` carries a long digit run (7+ digits, allowing the
/// separators a phone number, SSN, or card number is typically written
/// with — spaces, `-`, `.`, parens). Any non-digit, non-separator character
/// resets the run, so an ordinary sentence with a short number ("3 dogs")
/// never trips it.
fn has_long_digit_run(text: &str) -> bool {
    const MIN_RUN: usize = 7;
    let mut run = 0usize;
    for ch in text.chars() {
        if ch.is_ascii_digit() {
            run += 1;
            if run >= MIN_RUN {
                return true;
            }
        } else if matches!(ch, '-' | '.' | ' ' | '(' | ')' | '+') {
            // Separator: keep accumulating across it.
        } else {
            run = 0;
        }
    }
    false
}

/// The post-parse PII heuristic (`#796`): a refused-category keyword or a
/// long digit run anywhere in the text.
///
/// Runs on every candidate fact, every compaction-digest line, and every
/// deliberate memory-write note regardless of confidence — a confident PII
/// hit is exactly the dangerous case, not an exception to it.
///
/// Public because it is the ONE PII refusal predicate for durable memory:
/// fact extraction, the compaction digest pass (`#1149`), and the
/// deliberate `memory_write` path (`#1139`, INV-C9) all refuse through this
/// same function, so no write path can ever drift on what memory must never
/// carry.
#[must_use]
pub fn looks_like_pii(text: &str) -> bool {
    if has_long_digit_run(text) {
        return true;
    }
    let lower = text.to_lowercase();
    PII_REFUSAL_KEYWORDS.iter().any(|kw| lower.contains(kw))
}

/// Parse the model reply tolerantly: take the outermost `{…}` slice, decode,
/// drop empty/oversized entries, clamp confidence, apply the write-time
/// confidence floor and the PII heuristic (`#796`), and cap the batch. ANY
/// parse failure extracts nothing — never an error the pipeline would
/// propagate.
fn parse_reply(text: &str) -> ExtractedMemories {
    let Some(start) = text.find('{') else {
        return ExtractedMemories::default();
    };
    let Some(end) = text.rfind('}') else {
        return ExtractedMemories::default();
    };
    let Ok(wire) = serde_json::from_str::<WireReply>(&text[start..=end]) else {
        tracing::debug!("memory extractor reply was not the expected JSON; extracting nothing");
        return ExtractedMemories::default();
    };
    let added = wire
        .added
        .into_iter()
        .filter(|f| !f.text.trim().is_empty())
        .filter_map(|f| {
            let confidence_bps = f.confidence.min(100) * 100;
            if confidence_bps < MIN_CONFIDENCE_BPS {
                tracing::debug!(
                    confidence_bps,
                    floor = MIN_CONFIDENCE_BPS,
                    "extracted fact below the write-time confidence floor; dropped"
                );
                return None;
            }
            let text = f.text.trim().to_owned();
            if looks_like_pii(&text) {
                tracing::info!("extracted fact matched a PII refusal category; dropped (#796)");
                return None;
            }
            Some(CandidateFact {
                text,
                entities: f
                    .entities
                    .into_iter()
                    .filter(|e| !e.trim().is_empty())
                    .collect(),
                confidence_bps,
                replaces: {
                    let id = f.replaces.trim();
                    (!id.is_empty()).then(|| id.to_owned())
                },
            })
        })
        .take(MAX_FACTS_PER_TURN)
        .collect();
    let invalidated: Vec<Invalidation> = wire
        .invalidated
        .into_iter()
        .filter(|i| !i.fact_id.trim().is_empty())
        .map(|i| Invalidation {
            fact_id: i.fact_id.trim().to_owned(),
            reason: if i.reason.trim().is_empty() {
                "contradicted".to_owned()
            } else {
                i.reason.trim().to_owned()
            },
        })
        .collect();
    // A restated fact is corroboration, a contradiction is invalidation — never
    // both. If the model listed an id under both, invalidation wins: the merge
    // pass must not reinforce a fact this same turn also contradicted.
    let invalidated_ids: std::collections::HashSet<&str> =
        invalidated.iter().map(|i| i.fact_id.as_str()).collect();
    let mut seen = std::collections::HashSet::new();
    let corroborated = wire
        .corroborated
        .into_iter()
        .filter_map(|id| {
            let id = id.trim();
            (!id.is_empty() && !invalidated_ids.contains(id) && seen.insert(id.to_owned()))
                .then(|| id.to_owned())
        })
        .collect();
    ExtractedMemories {
        added,
        invalidated,
        corroborated,
    }
}

/// Distill one committed turn into memory operations.
///
/// Builds a [`CompletionRequest`] for `model` carrying the extraction
/// instructions, the persona's existing active facts, and the turn
/// transcript; runs it through `provider`; and parses the JSON reply
/// tolerantly (an unparseable reply extracts nothing). An empty `model`
/// defers to the provider's configured default — the right call for a
/// dedicated classifier provider; a non-empty value overrides it
/// per-request. Keep whichever applies pointed at a fast, inexpensive
/// backend — this runs after every committed turn.
///
/// # Errors
///
/// Propagates `P::Error` from [`LlmProvider::complete`] (pre-stream
/// failures) and from [`collect_turn`] (mid-stream faults) — the caller
/// logs and drops; extraction is best-effort by design.
pub async fn extract_memories<P: LlmProvider + ?Sized>(
    provider: &P,
    model: &str,
    transcript: &[ParticipationMsg],
    existing: &[ExistingFact],
) -> Result<ExtractedMemories, P::Error> {
    let mut req = CompletionRequest::new(model);
    req.messages.push(Message {
        role: Role::System,
        content: vec![Content::Text(system_prompt().to_owned())],
    });
    req.messages.push(Message {
        role: Role::User,
        content: vec![Content::Text(render_input(transcript, existing))],
    });
    let stream = provider.complete(req).await?;
    let out = collect_turn(stream).await?;
    Ok(parse_reply(&out.text))
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use std::sync::{Arc, Mutex};

    use async_trait::async_trait;
    use futures::stream::{self, BoxStream, StreamExt};
    use polyc_llm::{Chunk, StopReason, error::DummyError};

    use super::*;

    #[derive(Clone)]
    struct MockProvider {
        reply: String,
        captured: Arc<Mutex<Option<CompletionRequest>>>,
    }

    impl MockProvider {
        fn new(reply: &str) -> Self {
            Self {
                reply: reply.to_owned(),
                captured: Arc::new(Mutex::new(None)),
            }
        }
    }

    #[async_trait]
    impl LlmProvider for MockProvider {
        type Error = DummyError;

        async fn complete(
            &self,
            req: CompletionRequest,
        ) -> Result<BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error> {
            *self.captured.lock().unwrap() = Some(req);
            let chunks = vec![
                Ok(Chunk::text_delta(self.reply.clone())),
                Ok(Chunk::Stop(StopReason::EndTurn)),
            ];
            Ok(stream::iter(chunks).boxed())
        }
    }

    fn transcript() -> Vec<ParticipationMsg> {
        vec![
            ParticipationMsg {
                speaker: "erica".to_owned(),
                text: "actually I've switched to filter coffee".to_owned(),
                is_self: false,
            },
            ParticipationMsg {
                speaker: "bot".to_owned(),
                text: "noted!".to_owned(),
                is_self: true,
            },
        ]
    }

    #[tokio::test]
    async fn well_formed_reply_parses_adds_and_invalidations() {
        let provider = MockProvider::new(
            r#"{"added":[{"text":"prefers filter coffee","entities":["coffee"],"confidence":90,"replaces":"f1"}],
                "invalidated":[{"fact_id":"f1","reason":"switched"}]}"#,
        );
        let existing = [ExistingFact {
            fact_id: "f1".to_owned(),
            text: "prefers espresso".to_owned(),
        }];
        let out = extract_memories(&provider, "fast", &transcript(), &existing)
            .await
            .expect("extract");
        assert_eq!(out.added.len(), 1);
        assert_eq!(out.added[0].text, "prefers filter coffee");
        assert_eq!(out.added[0].confidence_bps, 9_000);
        assert_eq!(
            out.added[0].replaces.as_deref(),
            Some("f1"),
            "the replacement pairing survives parsing"
        );
        assert_eq!(out.invalidated.len(), 1);
        assert_eq!(out.invalidated[0].fact_id, "f1");
    }

    /// The `corroborated` restatement signal (#860) parses, dedups, and never
    /// overlaps with `invalidated`: a fact the same turn contradicted is never
    /// also reinforced (invalidation wins).
    #[tokio::test]
    async fn corroborated_ids_parse_dedup_and_exclude_contradictions() {
        let provider = MockProvider::new(
            r#"{"added":[],
                "invalidated":[{"fact_id":"f2","reason":"changed"}],
                "corroborated":["f1"," f1 ","  ","f2"]}"#,
        );
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(
            out.corroborated,
            vec!["f1".to_owned()],
            "f1 dedups to one; blanks drop; f2 is excluded (it was invalidated)"
        );
    }

    #[tokio::test]
    async fn missing_or_blank_replaces_parses_as_none() {
        let provider = MockProvider::new(
            r#"{"added":[{"text":"works UTC+2","confidence":80},
                          {"text":"has a dog","confidence":70,"replaces":"  "}]}"#,
        );
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(out.added.len(), 2);
        assert!(out.added.iter().all(|f| f.replaces.is_none()));
    }

    #[tokio::test]
    async fn prose_wrapped_json_still_parses() {
        let provider = MockProvider::new(
            "Here you go:\n{\"added\":[{\"text\":\"works UTC+2\",\"confidence\":80}],\"invalidated\":[]}\nDone.",
        );
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(out.added.len(), 1);
        assert_eq!(out.added[0].confidence_bps, 8_000);
    }

    #[tokio::test]
    async fn garbage_reply_extracts_nothing() {
        let provider = MockProvider::new("no json here at all");
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(out, ExtractedMemories::default());
    }

    #[tokio::test]
    async fn malformed_json_extracts_nothing() {
        let provider = MockProvider::new(r#"{"added": [{"text": 12}], "invalid"#);
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(out, ExtractedMemories::default());
    }

    #[tokio::test]
    async fn empty_texts_and_over_cap_batches_are_bounded() {
        let many: Vec<String> = (0..20)
            .map(|i| format!(r#"{{"text":"fact {i}","confidence":300}}"#))
            .collect();
        let provider = MockProvider::new(&format!(
            r#"{{"added":[{},{}],"invalidated":[{{"fact_id":"  "}}]}}"#,
            r#"{"text":"   "}"#,
            many.join(",")
        ));
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(out.added.len(), MAX_FACTS_PER_TURN, "batch is capped");
        assert!(
            out.added.iter().all(|f| f.confidence_bps <= 10_000),
            "confidence clamps to 100%"
        );
        assert!(
            out.invalidated.is_empty(),
            "blank fact ids are dropped, not passed through"
        );
    }

    /// Write-time confidence floor (`#796`, defect #2): a low-confidence fact
    /// is dropped at parse time even though it is otherwise well-formed — the
    /// extractor never gets to persist something it wasn't sure of.
    #[tokio::test]
    async fn low_confidence_fact_is_dropped() {
        let provider = MockProvider::new(
            r#"{"added":[
                {"text":"maybe prefers tea, not certain","confidence":40},
                {"text":"definitely prefers filter coffee","confidence":95}
            ]}"#,
        );
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(out.added.len(), 1, "the below-floor fact is dropped");
        assert_eq!(out.added[0].text, "definitely prefers filter coffee");
    }

    /// The PII heuristic (`#796`, defect #2): a private-address or
    /// health/credential fact is refused regardless of confidence — a
    /// confident PII fact is the dangerous case, not an exception.
    #[tokio::test]
    async fn pii_facts_are_refused_even_at_high_confidence() {
        let provider = MockProvider::new(
            r#"{"added":[
                {"text":"home address is 42 Rowan Street","confidence":99},
                {"text":"was diagnosed with a chronic condition","confidence":99},
                {"text":"phone number is 555-123-4567","confidence":99},
                {"text":"prefers filter coffee","confidence":99}
            ]}"#,
        );
        let out = extract_memories(&provider, "fast", &transcript(), &[])
            .await
            .expect("extract");
        assert_eq!(
            out.added.len(),
            1,
            "only the non-PII fact survives: {:?}",
            out.added
        );
        assert_eq!(out.added[0].text, "prefers filter coffee");
    }

    #[tokio::test]
    async fn request_carries_existing_facts_and_transcript() {
        let provider = MockProvider::new("{}");
        let existing = [ExistingFact {
            fact_id: "f1".to_owned(),
            text: "prefers espresso".to_owned(),
        }];
        let _ = extract_memories(&provider, "fast", &transcript(), &existing)
            .await
            .expect("extract");
        let req = provider.captured.lock().unwrap().clone().expect("captured");
        assert_eq!(req.messages.len(), 2);
        let user_text = match &req.messages[1].content[0] {
            Content::Text(t) => t.clone(),
            other => panic!("expected text, got {other:?}"),
        };
        assert!(user_text.contains("[f1] prefers espresso"));
        assert!(user_text.contains("erica: actually I've switched"));
        assert!(user_text.contains("assistant: noted!"));
    }
}