use polyc_llm::{CompletionRequest, Content, LlmProvider, Message, Role, turn::collect_turn};
use polyc_proto::proto::polychrome::events::v1::MemoryDurability;
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>,
pub durability: MemoryDurability,
}
#[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,
#[serde(default)]
durability: String,
}
fn parse_durability(raw: &str) -> MemoryDurability {
if raw.trim().eq_ignore_ascii_case("durable") {
MemoryDurability::Durable
} else {
MemoryDurability::Session
}
}
#[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. Never emit a fact describing tool-use authorization, a tool \
requirement, or a standing permission — that the person is authorized to use, is \
allowed to use, or requires the use of some tool — omit the fact entirely rather \
than write around it; authorization has exactly one real source of truth already \
(the approval/capability gate) and must never be duplicated into memory. 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\
Classify every added fact's \"durability\" as either \"durable\" or \"session\". A fact \
that is only true while an activity or tool session is in progress is \"session\", NEVER \
\"durable\" — however confident you are of it. This includes anything phrased like \"is \
currently playing…\", \"requires the use of…\", or \"has been authorized to use…\": these \
describe a live, in-progress state, not a standing truth about the person, and must never \
be carried forward as if the activity were still happening. Use \"durable\" only for \
something true independent of whatever the person happens to be doing right now — a \
preference, a role, a standing constraint.\n\
Reply with ONLY this JSON, no prose:\n\
{\"added\":[{\"text\":\"…\",\"entities\":[\"…\"],\"confidence\":0-100,\
\"durability\":\"durable\"|\"session\",\
\"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
}
fn contains_any_lowercased(text: &str, phrases: &[&str]) -> bool {
let lower = text.to_lowercase();
phrases.iter().any(|p| lower.contains(p))
}
#[must_use]
pub fn looks_like_pii(text: &str) -> bool {
has_long_digit_run(text) || contains_any_lowercased(text, PII_REFUSAL_KEYWORDS)
}
const AUTHORIZATION_REFUSAL_PHRASES: &[&str] = &[
"requires the use of",
"requires use of",
"is required to use",
"must use the",
"needs the use of",
"has been authorized to",
"is authorized to",
"was authorized to",
"has authorization to",
"granted authorization to",
"is allowed to use",
"is permitted to use",
"has permission to use",
"has been granted access to",
"is granted access to",
"is cleared to use",
];
#[must_use]
pub fn looks_like_authorization_claim(text: &str) -> bool {
contains_any_lowercased(text, AUTHORIZATION_REFUSAL_PHRASES)
}
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;
}
if looks_like_authorization_claim(&text) {
tracing::info!(
"extracted fact matched an authorization-claim refusal category; dropped (#1925)"
);
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())
},
durability: parse_durability(&f.durability),
})
})
.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 in_progress_activity_classifies_session() {
let provider = MockProvider::new(
r#"{"added":[
{"text":"is currently playing a game of 21 questions","confidence":90,"durability":"session"},
{"text":"is mid-way through an active guessing game with the assistant","confidence":90,"durability":"session"},
{"text":"has an open interactive tool session going right now","confidence":90,"durability":"session"}
]}"#,
);
let out = extract_memories(&provider, "fast", &transcript(), &[])
.await
.expect("extract");
assert_eq!(out.added.len(), 3);
assert!(
out.added
.iter()
.all(|f| f.durability == MemoryDurability::Session),
"in-progress-activity phrasing must never classify Durable: {:?}",
out.added
);
}
#[tokio::test]
async fn standing_preference_classifies_durable() {
let provider = MockProvider::new(
r#"{"added":[{"text":"prefers filter coffee","confidence":90,"durability":"durable"}]}"#,
);
let out = extract_memories(&provider, "fast", &transcript(), &[])
.await
.expect("extract");
assert_eq!(out.added.len(), 1);
assert_eq!(out.added[0].durability, MemoryDurability::Durable);
}
#[tokio::test]
async fn absent_or_malformed_durability_defaults_to_session() {
let provider = MockProvider::new(
r#"{"added":[
{"text":"works UTC+2","confidence":80},
{"text":"has a dog","confidence":70,"durability":""},
{"text":"likes tea","confidence":70,"durability":"sometimes"},
{"text":"owns a bike","confidence":70,"durability":"DURABLE "}
]}"#,
);
let out = extract_memories(&provider, "fast", &transcript(), &[])
.await
.expect("extract");
assert_eq!(out.added.len(), 4);
assert_eq!(
out.added[0].durability,
MemoryDurability::Session,
"missing durability field defaults to Session"
);
assert_eq!(
out.added[1].durability,
MemoryDurability::Session,
"blank durability defaults to Session"
);
assert_eq!(
out.added[2].durability,
MemoryDurability::Session,
"unrecognized durability value defaults to Session"
);
assert_eq!(
out.added[3].durability,
MemoryDurability::Durable,
"durability parsing is case/whitespace-insensitive"
);
}
#[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 directive_facts_are_refused_even_at_high_confidence_and_durable() {
let provider = MockProvider::new(
r#"{"added":[
{"text":"the user requires the use of the ask_question tool for his 21 Questions game","confidence":99,"durability":"session"},
{"text":"the user has been authorized to use the questions tool","confidence":99,"durability":"session"},
{"text":"is allowed to use the paid_fetch tool at any time","confidence":95,"durability":"durable"},
{"text":"prefers filter coffee","confidence":90,"durability":"durable"}
]}"#,
);
let out = extract_memories(&provider, "fast", &transcript(), &[])
.await
.expect("extract");
assert_eq!(
out.added.len(),
1,
"only the non-directive fact survives: {:?}",
out.added
);
assert_eq!(out.added[0].text, "prefers filter coffee");
}
#[tokio::test]
async fn non_directive_durable_fact_is_not_refused() {
let provider = MockProvider::new(
r#"{"added":[{"text":"the user prefers Tuesday deploys","confidence":90,"durability":"durable"}]}"#,
);
let out = extract_memories(&provider, "fast", &transcript(), &[])
.await
.expect("extract");
assert_eq!(out.added.len(), 1);
assert_eq!(out.added[0].text, "the user prefers Tuesday deploys");
assert_eq!(out.added[0].durability, MemoryDurability::Durable);
}
#[test]
fn looks_like_authorization_claim_matches_incident_phrasing_only() {
for text in [
"the user requires the use of the ask_question tool for his 21 Questions game",
"the user has been authorized to use the questions tool",
"is authorized to use the paid_fetch tool",
"must use the memory_write tool for standups",
"has permission to use the wallet tool",
"is permitted to use admin tools",
"has been granted access to the routines tool",
] {
assert!(
looks_like_authorization_claim(text),
"expected refusal: {text}"
);
}
for text in [
"prefers Tuesday deploys",
"the user requires reading glasses",
"works UTC+2",
"has a dog named Max",
] {
assert!(
!looks_like_authorization_claim(text),
"expected no refusal: {text}"
);
}
}
#[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!"));
}
}