use polyc_llm::{CompletionRequest, Content, LlmProvider, Message, Role, turn::collect_turn};
use serde::Deserialize;
use crate::participation::ParticipationMsg;
pub const MAX_FACTS_PER_TURN: usize = 8;
pub const MIN_CONFIDENCE_BPS: u32 = 6_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CandidateFact {
pub text: String,
pub entities: Vec<String>,
pub confidence_bps: u32,
pub replaces: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Invalidation {
pub fact_id: String,
pub reason: String,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ExtractedMemories {
pub added: Vec<CandidateFact>,
pub invalidated: Vec<Invalidation>,
pub corroborated: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct ExistingFact {
pub fact_id: String,
pub text: String,
}
#[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>,
#[serde(default)]
confidence: u32,
#[serde(default)]
replaces: String,
}
#[derive(Debug, Deserialize)]
struct WireInvalidation {
#[serde(default)]
fact_id: String,
#[serde(default)]
reason: String,
}
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."
}
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
}
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",
];
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, '-' | '.' | ' ' | '(' | ')' | '+') {
} else {
run = 0;
}
}
false
}
#[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))
}
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();
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,
}
}
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 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"
);
}
#[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");
}
#[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!"));
}
}