use chrono::Utc; use opensearch::OpenSearch; use serde_json::json; use uuid::Uuid; use super::schema::MemoryDocument; /// Search memories by content relevance, filtered to a specific user. pub async fn query( client: &OpenSearch, index: &str, user_id: &str, query_text: &str, limit: usize, ) -> anyhow::Result> { let body = json!({ "size": limit, "query": { "bool": { "filter": [ { "term": { "user_id": user_id } } ], "must": [ { "match": { "content": query_text } } ] } }, "sort": [{ "_score": "desc" }] }); let response = client .search(opensearch::SearchParts::Index(&[index])) .body(body) .send() .await?; let data: serde_json::Value = response.json().await?; parse_hits(&data) } /// Get the most recent memories for a user, sorted by updated_at desc. pub async fn get_recent( client: &OpenSearch, index: &str, user_id: &str, limit: usize, ) -> anyhow::Result> { let body = json!({ "size": limit, "query": { "bool": { "filter": [ { "term": { "user_id": user_id } } ] } }, "sort": [{ "updated_at": "desc" }] }); let response = client .search(opensearch::SearchParts::Index(&[index])) .body(body) .send() .await?; let data: serde_json::Value = response.json().await?; parse_hits(&data) } /// Store a new memory document for a user. pub async fn set( client: &OpenSearch, index: &str, user_id: &str, content: &str, category: &str, source: &str, ) -> anyhow::Result<()> { let now = Utc::now().timestamp_millis(); let id = Uuid::new_v4().to_string(); let doc = MemoryDocument { id: id.clone(), user_id: user_id.to_string(), content: content.to_string(), category: category.to_string(), created_at: now, updated_at: now, source: source.to_string(), }; let response = client .index(opensearch::IndexParts::IndexId(index, &id)) .body(serde_json::to_value(&doc)?) .send() .await?; if !response.status_code().is_success() { let body = response.text().await?; anyhow::bail!("Failed to store memory: {body}"); } Ok(()) } pub(crate) fn parse_hits(data: &serde_json::Value) -> anyhow::Result> { let hits = data["hits"]["hits"] .as_array() .cloned() .unwrap_or_default(); let mut docs = Vec::with_capacity(hits.len()); for hit in &hits { if let Ok(doc) = serde_json::from_value::(hit["_source"].clone()) { docs.push(doc); } } Ok(docs) } #[cfg(test)] mod tests { use super::*; fn fake_os_response(sources: Vec) -> serde_json::Value { let hits: Vec = sources .into_iter() .map(|s| json!({ "_source": s })) .collect(); json!({ "hits": { "hits": hits } }) } #[test] fn test_parse_hits_multiple() { let data = fake_os_response(vec![ json!({ "id": "a", "user_id": "sienna@sunbeam.pt", "content": "prefers terse answers", "category": "preference", "created_at": 1710000000000_i64, "updated_at": 1710000000000_i64, "source": "auto" }), json!({ "id": "b", "user_id": "sienna@sunbeam.pt", "content": "working on drive UI", "category": "fact", "created_at": 1710000000000_i64, "updated_at": 1710000000000_i64, "source": "script" }), ]); let docs = parse_hits(&data).unwrap(); assert_eq!(docs.len(), 2); assert_eq!(docs[0].id, "a"); assert_eq!(docs[0].content, "prefers terse answers"); assert_eq!(docs[1].id, "b"); assert_eq!(docs[1].category, "fact"); } #[test] fn test_parse_hits_empty() { let data = json!({ "hits": { "hits": [] } }); let docs = parse_hits(&data).unwrap(); assert!(docs.is_empty()); } #[test] fn test_parse_hits_missing_structure() { let data = json!({}); let docs = parse_hits(&data).unwrap(); assert!(docs.is_empty()); } #[test] fn test_parse_hits_skips_malformed() { let data = fake_os_response(vec![ json!({ "id": "good", "user_id": "x@y", "content": "ok", "category": "fact", "created_at": 1, "updated_at": 1, "source": "auto" }), json!({ "bad": "no fields" }), ]); let docs = parse_hits(&data).unwrap(); assert_eq!(docs.len(), 1); assert_eq!(docs[0].id, "good"); } }