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