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;

/// 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>,
}

/// 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>,
}

#[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. 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.\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\":\"…\"}]}\n\
     Both 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
}

/// Parse the model reply tolerantly: take the outermost `{…}` slice, decode,
/// drop empty/oversized entries, clamp confidence, 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())
        .take(MAX_FACTS_PER_TURN)
        .map(|f| CandidateFact {
            text: f.text.trim().to_owned(),
            entities: f
                .entities
                .into_iter()
                .filter(|e| !e.trim().is_empty())
                .collect(),
            confidence_bps: f.confidence.min(100) * 100,
            replaces: {
                let id = f.replaces.trim();
                (!id.is_empty()).then(|| id.to_owned())
            },
        })
        .collect();
    let invalidated = 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();
    ExtractedMemories { added, invalidated }
}

/// 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");
    }

    #[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"
        );
    }

    #[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!"));
    }
}