1use zeph_llm::provider::{LlmProvider as _, Message, MessageMetadata, Role};
5
6use super::{KEY_FACTS_COLLECTION, SemanticMemory};
7use crate::embedding_store::MessageKind;
8use crate::error::MemoryError;
9use crate::types::{ConversationId, MessageId};
10use crate::vector_store::{FieldCondition, FieldValue, VectorFilter};
11
12#[derive(Debug, Clone, serde::Deserialize, serde::Serialize, schemars::JsonSchema)]
13pub struct StructuredSummary {
14 pub summary: String,
15 pub key_facts: Vec<String>,
16 pub entities: Vec<String>,
17}
18
19#[derive(Debug, Clone)]
20pub struct Summary {
21 pub id: i64,
22 pub conversation_id: ConversationId,
23 pub content: String,
24 pub first_message_id: Option<MessageId>,
26 pub last_message_id: Option<MessageId>,
28 pub token_estimate: i64,
29}
30
31#[derive(Debug, Clone, Copy)]
33pub struct SummarizeOutcome {
34 pub summary_id: i64,
36 pub messages_folded: usize,
39}
40
41#[must_use]
42pub fn build_summarization_prompt(messages: &[(MessageId, String, String)]) -> String {
43 let mut prompt = String::from(
44 "Summarize the following conversation. Extract key facts, decisions, entities, \
45 and context needed to continue the conversation.\n\n\
46 Respond in JSON with fields: summary (string), key_facts (list of strings), \
47 entities (list of strings).\n\nConversation:\n",
48 );
49
50 for (_, role, content) in messages {
51 prompt.push_str(role);
52 prompt.push_str(": ");
53 prompt.push_str(content);
54 prompt.push('\n');
55 }
56
57 prompt
58}
59
60impl SemanticMemory {
61 pub async fn load_summaries(
67 &self,
68 conversation_id: ConversationId,
69 ) -> Result<Vec<Summary>, MemoryError> {
70 let rows = self.sqlite.load_summaries(conversation_id).await?;
71 let summaries = rows
72 .into_iter()
73 .map(
74 |(
75 id,
76 conversation_id,
77 content,
78 first_message_id,
79 last_message_id,
80 token_estimate,
81 )| {
82 Summary {
83 id,
84 conversation_id,
85 content,
86 first_message_id,
87 last_message_id,
88 token_estimate,
89 }
90 },
91 )
92 .collect();
93 Ok(summaries)
94 }
95
96 #[tracing::instrument(name = "memory.summarize", skip_all, fields(input_msgs = %message_count, output_len = tracing::field::Empty))]
107 pub async fn summarize(
108 &self,
109 conversation_id: ConversationId,
110 message_count: usize,
111 ) -> Result<Option<SummarizeOutcome>, MemoryError> {
112 let total = self.sqlite.count_messages(conversation_id).await?;
113
114 if total <= i64::try_from(message_count)? {
115 return Ok(None);
116 }
117
118 let after_id = self
119 .sqlite
120 .latest_summary_last_message_id(conversation_id)
121 .await?
122 .unwrap_or(MessageId(0));
123
124 let messages = self
125 .sqlite
126 .load_messages_range(conversation_id, after_id, message_count)
127 .await?;
128
129 if messages.is_empty() {
130 return Ok(None);
131 }
132
133 let messages_folded = messages.len();
134 let prompt = build_summarization_prompt(&messages);
135 let chat_messages = vec![Message {
136 role: Role::User,
137 content: prompt,
138 parts: vec![],
139 metadata: MessageMetadata::default(),
140 }];
141
142 let structured = self.call_summarization_llm(&chat_messages).await?;
143 let summary_text = &structured.summary;
144
145 let token_estimate = i64::try_from(self.token_counter.count_tokens(summary_text))?;
146 let first_message_id = messages[0].0;
147 let last_message_id = messages[messages.len() - 1].0;
148
149 let summary_id = self
150 .sqlite
151 .save_summary(
152 conversation_id,
153 summary_text,
154 Some(first_message_id),
155 Some(last_message_id),
156 token_estimate,
157 )
158 .await?;
159
160 if let Some(qdrant) = &self.qdrant
161 && self.effective_embed_provider().supports_embeddings()
162 {
163 match tokio::time::timeout(
164 self.embed_timeout,
165 self.effective_embed_provider().embed(summary_text),
166 )
167 .await
168 {
169 Ok(Ok(vector)) => {
170 if let Err(e) = qdrant.ensure_collection_for_vector(&vector).await {
171 tracing::warn!("Failed to ensure Qdrant collection: {e:#}");
172 } else if let Err(e) = qdrant
173 .store(
174 MessageId(summary_id),
175 conversation_id,
176 "system",
177 vector,
178 MessageKind::Summary,
179 &self.embedding_model,
180 0,
181 )
182 .await
183 {
184 tracing::warn!("Failed to embed summary: {e:#}");
185 }
186 }
187 Ok(Err(e)) => {
188 tracing::warn!("Failed to generate summary embedding: {e:#}");
189 }
190 Err(_) => {
191 tracing::warn!("summarize: embed timed out for summary text — skipping store");
192 }
193 }
194 }
195
196 if !structured.key_facts.is_empty() {
197 self.store_key_facts(conversation_id, summary_id, &structured.key_facts)
198 .await;
199 }
200
201 Ok(Some(SummarizeOutcome {
202 summary_id,
203 messages_folded,
204 }))
205 }
206
207 async fn call_summarization_llm(
216 &self,
217 chat_messages: &[Message],
218 ) -> Result<StructuredSummary, MemoryError> {
219 let timeout_secs = self.summarization_llm_timeout_secs;
220 let timeout = std::time::Duration::from_secs(timeout_secs);
221 match tokio::time::timeout(
222 timeout,
223 self.provider
224 .chat_typed_erased::<StructuredSummary>(chat_messages),
225 )
226 .await
227 {
228 Ok(Ok(s)) => Ok(s),
229 Ok(Err(e)) => {
230 tracing::warn!(
231 "structured summarization failed, falling back to plain text: {e:#}"
232 );
233 match tokio::time::timeout(timeout, self.provider.chat(chat_messages)).await {
234 Ok(Ok(plain)) => Ok(StructuredSummary {
235 summary: plain,
236 key_facts: vec![],
237 entities: vec![],
238 }),
239 Ok(Err(e)) => Err(MemoryError::Llm(e)),
240 Err(_elapsed) => {
241 tracing::warn!(
242 "summarization: plain text fallback LLM call timed out after {timeout_secs}s"
243 );
244 Err(MemoryError::Timeout("LLM call timed out".into()))
245 }
246 }
247 }
248 Err(_elapsed) => {
249 tracing::warn!(
250 "summarization: structured LLM call timed out after {timeout_secs}s"
251 );
252 Err(MemoryError::Timeout("LLM call timed out".into()))
253 }
254 }
255 }
256
257 pub(super) async fn store_key_facts(
258 &self,
259 conversation_id: ConversationId,
260 source_summary_id: i64,
261 key_facts: &[String],
262 ) {
263 let Some(qdrant) = &self.qdrant else {
264 return;
265 };
266 if !self.effective_embed_provider().supports_embeddings() {
267 return;
268 }
269
270 let filtered: Vec<&str> = key_facts
274 .iter()
275 .filter(|f| !is_policy_decision_fact(f.as_str()))
276 .map(String::as_str)
277 .collect();
278
279 let Some(first_fact) = filtered.first().copied() else {
280 return;
281 };
282 let first_vector = match tokio::time::timeout(
283 self.embed_timeout,
284 self.effective_embed_provider().embed(first_fact),
285 )
286 .await
287 {
288 Ok(Ok(v)) => v,
289 Ok(Err(e)) => {
290 tracing::warn!("Failed to embed key fact: {e:#}");
291 return;
292 }
293 Err(_) => {
294 tracing::warn!("store_key_facts: embed timed out for first fact — skipping");
295 return;
296 }
297 };
298 if let Err(e) = qdrant
299 .ensure_named_collection_for_vector(KEY_FACTS_COLLECTION, &first_vector)
300 .await
301 {
302 tracing::warn!("Failed to ensure key_facts collection: {e:#}");
303 return;
304 }
305
306 let threshold = self.key_facts_dedup_threshold;
307 self.store_key_fact_if_unique(
308 qdrant,
309 conversation_id,
310 source_summary_id,
311 first_fact,
312 first_vector,
313 threshold,
314 )
315 .await;
316
317 for fact in filtered[1..].iter().copied() {
318 match tokio::time::timeout(
319 self.embed_timeout,
320 self.effective_embed_provider().embed(fact),
321 )
322 .await
323 {
324 Ok(Ok(vector)) => {
325 self.store_key_fact_if_unique(
326 qdrant,
327 conversation_id,
328 source_summary_id,
329 fact,
330 vector,
331 threshold,
332 )
333 .await;
334 }
335 Ok(Err(e)) => {
336 tracing::warn!("Failed to embed key fact: {e:#}");
337 }
338 Err(_) => {
339 tracing::warn!("store_key_facts: embed timed out for fact — skipping");
340 }
341 }
342 }
343 }
344
345 async fn store_key_fact_if_unique(
346 &self,
347 qdrant: &crate::embedding_store::EmbeddingStore,
348 conversation_id: ConversationId,
349 source_summary_id: i64,
350 fact: &str,
351 vector: Vec<f32>,
352 threshold: f32,
353 ) {
354 let dedup_filter = Some(VectorFilter {
359 must: vec![
360 FieldCondition {
361 field: "conversation_id".into(),
362 value: FieldValue::Integer(conversation_id.0),
363 },
364 FieldCondition {
365 field: "db_instance_id".into(),
366 value: FieldValue::Text(qdrant.db_instance_id().to_owned()),
367 },
368 ],
369 must_not: vec![],
370 });
371 match qdrant
372 .search_collection(KEY_FACTS_COLLECTION, &vector, 1, dedup_filter)
373 .await
374 {
375 Ok(hits) if hits.first().is_some_and(|h| h.score >= threshold) => {
376 tracing::debug!(
377 score = hits[0].score,
378 threshold,
379 "key-facts: skipping near-duplicate fact"
380 );
381 return;
382 }
383 Ok(_) => {}
384 Err(e) => {
385 tracing::warn!("key-facts: dedup search failed, storing anyway: {e:#}");
386 }
387 }
388
389 let payload = serde_json::json!({
390 "conversation_id": conversation_id.0,
391 "db_instance_id": qdrant.db_instance_id(),
392 "fact_text": fact,
393 "source_summary_id": source_summary_id,
394 });
395 if let Err(e) = qdrant
396 .store_to_collection(KEY_FACTS_COLLECTION, payload, vector)
397 .await
398 {
399 tracing::warn!("Failed to store key fact: {e:#}");
400 }
401 }
402
403 pub async fn search_key_facts(
414 &self,
415 query: &str,
416 limit: usize,
417 conversation_id: Option<ConversationId>,
418 ) -> Result<Vec<String>, MemoryError> {
419 let Some(qdrant) = &self.qdrant else {
420 tracing::debug!("key-facts: skipped, no vector store");
421 return Ok(Vec::new());
422 };
423 if !self.effective_embed_provider().supports_embeddings() {
424 tracing::debug!("key-facts: skipped, no embedding support");
425 return Ok(Vec::new());
426 }
427
428 let vector = match tokio::time::timeout(
429 self.embed_timeout,
430 self.effective_embed_provider().embed(query),
431 )
432 .await
433 {
434 Ok(Ok(v)) => v,
435 Ok(Err(e)) => return Err(e.into()),
436 Err(_) => {
437 tracing::warn!("search_key_facts: embed timed out, returning empty results");
438 return Ok(Vec::new());
439 }
440 };
441 qdrant
442 .ensure_named_collection_for_vector(KEY_FACTS_COLLECTION, &vector)
443 .await?;
444
445 let filter = conversation_id.map(|cid| VectorFilter {
446 must: vec![
447 FieldCondition {
448 field: "conversation_id".into(),
449 value: FieldValue::Integer(cid.0),
450 },
451 FieldCondition {
452 field: "db_instance_id".into(),
453 value: FieldValue::Text(qdrant.db_instance_id().to_owned()),
454 },
455 ],
456 must_not: vec![],
457 });
458
459 let points = qdrant
460 .search_collection(KEY_FACTS_COLLECTION, &vector, limit, filter)
461 .await?;
462
463 tracing::debug!(
464 results = points.len(),
465 limit,
466 conversation_id = conversation_id.map(|c| c.0),
467 "key-facts: search complete"
468 );
469
470 let facts = points
471 .into_iter()
472 .filter_map(|p| p.payload.get("fact_text")?.as_str().map(String::from))
473 .collect();
474
475 Ok(facts)
476 }
477
478 pub async fn search_document_collection(
488 &self,
489 collection: &str,
490 query: &str,
491 limit: usize,
492 ) -> Result<Vec<crate::ScoredVectorPoint>, MemoryError> {
493 let Some(qdrant) = &self.qdrant else {
494 return Ok(Vec::new());
495 };
496 if !self.effective_embed_provider().supports_embeddings() {
497 return Ok(Vec::new());
498 }
499 if !qdrant.collection_exists(collection).await? {
500 return Ok(Vec::new());
501 }
502 let vector = match tokio::time::timeout(
503 self.embed_timeout,
504 self.effective_embed_provider().embed(query),
505 )
506 .await
507 {
508 Ok(Ok(v)) => v,
509 Ok(Err(e)) => return Err(e.into()),
510 Err(_) => {
511 tracing::warn!(
512 "search_document_collection: embed timed out, returning empty results"
513 );
514 return Ok(Vec::new());
515 }
516 };
517 let results = qdrant
518 .search_collection(collection, &vector, limit, None)
519 .await?;
520
521 tracing::debug!(
522 results = results.len(),
523 limit,
524 collection,
525 "document-collection: search complete"
526 );
527
528 Ok(results)
529 }
530}
531
532pub(crate) fn is_policy_decision_fact(fact: &str) -> bool {
538 const MARKERS: &[&str] = &[
539 "blocked",
540 "skipped",
541 "cannot access",
542 "security polic",
543 "utility polic",
544 "not allowed",
545 "permission denied",
546 "access denied",
547 "was denied",
548 ];
549 let lower = fact.to_lowercase();
550 MARKERS.iter().any(|m| lower.contains(m))
551}