diff --git a/.gitignore b/.gitignore index 9933adf6ad1..2a185ac8bbc 100644 --- a/.gitignore +++ b/.gitignore @@ -66,6 +66,7 @@ tauri.key tauri.key.pub /target/ src-tauri/target/ +.target-codex/ workflow .fastembed_cache diff --git a/Cargo.lock b/Cargo.lock index 170f8578aeb..d48d5a32465 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4541,7 +4541,7 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" [[package]] name = "openhuman" -version = "0.53.3" +version = "0.53.4" dependencies = [ "aes-gcm", "anyhow", diff --git a/docs/memory-tree-async-pipeline.excalidraw b/docs/memory-tree-async-pipeline.excalidraw new file mode 100644 index 00000000000..29ee1a1ce55 --- /dev/null +++ b/docs/memory-tree-async-pipeline.excalidraw @@ -0,0 +1,1353 @@ +{ + "type": "excalidraw", + "version": 2, + "source": "openhuman-memory-tree-async-pipeline", + "elements": [ + { + "id": "title", + "type": "text", + "x": 355, + "y": 20, + "width": 740, + "height": 40, + "text": "OpenHuman Memory Tree Async Pipeline", + "fontSize": 30, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 32, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 1, + "version": 1, + "versionNonce": 1, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "subtitle", + "type": "text", + "x": 235, + "y": 64, + "width": 980, + "height": 24, + "text": "Leaf ingestion -> jobs queue -> workers -> source/topic/global tree building", + "fontSize": 18, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 18, + "strokeColor": "#495057", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 2, + "version": 1, + "versionNonce": 2, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane1", + "type": "rectangle", + "x": 40, + "y": 120, + "width": 310, + "height": 340, + "strokeColor": "#1971c2", + "backgroundColor": "#e7f5ff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 3, + "version": 1, + "versionNonce": 3, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane2", + "type": "rectangle", + "x": 390, + "y": 120, + "width": 340, + "height": 340, + "strokeColor": "#2b8a3e", + "backgroundColor": "#ebfbee", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 4, + "version": 1, + "versionNonce": 4, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane3", + "type": "rectangle", + "x": 770, + "y": 120, + "width": 390, + "height": 340, + "strokeColor": "#e67700", + "backgroundColor": "#fff4e6", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 5, + "version": 1, + "versionNonce": 5, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane4", + "type": "rectangle", + "x": 1200, + "y": 120, + "width": 440, + "height": 340, + "strokeColor": "#9c36b5", + "backgroundColor": "#f8f0fc", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 6, + "version": 1, + "versionNonce": 6, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane5", + "type": "rectangle", + "x": 40, + "y": 500, + "width": 760, + "height": 220, + "strokeColor": "#0b7285", + "backgroundColor": "#e3fafc", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 7, + "version": 1, + "versionNonce": 7, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "lane6", + "type": "rectangle", + "x": 840, + "y": 500, + "width": 800, + "height": 220, + "strokeColor": "#495057", + "backgroundColor": "#f1f3f5", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 8, + "version": 1, + "versionNonce": 8, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h1", + "type": "text", + "x": 135, + "y": 135, + "width": 120, + "height": 28, + "text": "1. Ingest", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 9, + "version": 1, + "versionNonce": 9, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h2", + "type": "text", + "x": 505, + "y": 135, + "width": 110, + "height": 28, + "text": "2. Queue", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#2b8a3e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 10, + "version": 1, + "versionNonce": 10, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h3", + "type": "text", + "x": 890, + "y": 135, + "width": 150, + "height": 28, + "text": "3. Workers", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#e67700", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 11, + "version": 1, + "versionNonce": 11, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h4", + "type": "text", + "x": 1320, + "y": 135, + "width": 200, + "height": 28, + "text": "4. Tree State", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#9c36b5", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 12, + "version": 1, + "versionNonce": 12, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h5", + "type": "text", + "x": 275, + "y": 515, + "width": 290, + "height": 28, + "text": "5. Scheduler / Background", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#0b7285", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 13, + "version": 1, + "versionNonce": 13, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "h6", + "type": "text", + "x": 1100, + "y": 515, + "width": 280, + "height": 28, + "text": "6. Leaf Lifecycle", + "fontSize": 24, + "fontFamily": 1, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 24, + "strokeColor": "#343a40", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 14, + "version": 1, + "versionNonce": 14, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b1", + "type": "rectangle", + "x": 75, + "y": 185, + "width": 240, + "height": 240, + "strokeColor": "#1971c2", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 15, + "version": 1, + "versionNonce": 15, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t1", + "type": "text", + "x": 93, + "y": 205, + "width": 204, + "height": 198, + "text": "JSON-RPC / source ingestion\n\nchat | email | document\n\ncanonicalise\n-> chunk_markdown\n-> score_chunks_fast\n-> upsert_chunks_tx\n-> lifecycle_status = pending_extraction\n-> persist fast score rows\n-> enqueue extract_chunk per chunk\n-> wake_workers()", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 194, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 16, + "version": 1, + "versionNonce": 16, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b2", + "type": "rectangle", + "x": 435, + "y": 185, + "width": 250, + "height": 240, + "strokeColor": "#2b8a3e", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 17, + "version": 1, + "versionNonce": 17, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t2", + "type": "text", + "x": 453, + "y": 205, + "width": 214, + "height": 198, + "text": "SQLite: memory_tree/chunks.db\n\nmem_tree_chunks\nmem_tree_score\nmem_tree_entity_index\nmem_tree_jobs\nmem_tree_trees\nmem_tree_buffers\nmem_tree_summaries\n\njobs fields\nkind | payload_json | dedupe_key\nstatus | attempts | available_at_ms\nlocked_until_ms | last_error", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 194, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 18, + "version": 1, + "versionNonce": 18, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b3", + "type": "rectangle", + "x": 815, + "y": 185, + "width": 300, + "height": 135, + "strokeColor": "#e67700", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 19, + "version": 1, + "versionNonce": 19, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t3", + "type": "text", + "x": 833, + "y": 205, + "width": 264, + "height": 108, + "text": "jobs::start(workspace_dir)\n\nrecover_stale_locks()\nspawn 3 worker tasks\nNotify wakeup + 5s polling fallback\nshared Semaphore(3) for LLM-bound work", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 104, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 20, + "version": 1, + "versionNonce": 20, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b4", + "type": "rectangle", + "x": 815, + "y": 335, + "width": 300, + "height": 90, + "strokeColor": "#d9480f", + "backgroundColor": "#fff8f0", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 21, + "version": 1, + "versionNonce": 21, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t4", + "type": "text", + "x": 833, + "y": 355, + "width": 264, + "height": 54, + "text": "Handlers\nextract_chunk | append_buffer | seal\ntopic_route | digest_daily | flush_stale", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 50, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 22, + "version": 1, + "versionNonce": 22, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b5", + "type": "rectangle", + "x": 1240, + "y": 185, + "width": 360, + "height": 240, + "strokeColor": "#9c36b5", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 23, + "version": 1, + "versionNonce": 23, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t5", + "type": "text", + "x": 1258, + "y": 205, + "width": 324, + "height": 198, + "text": "Tree building outputs\n\nsource tree\nL0 buffer -> seal -> L1/L2/... summaries\n\ntopic tree\ncurator hotness gate\noptional append_buffer(topic)\n\nglobal tree\ndigest_daily -> daily node\nappend_daily_and_cascade", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 194, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 24, + "version": 1, + "versionNonce": 24, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "b6", + "type": "rectangle", + "x": 85, + "y": 575, + "width": 670, + "height": 105, + "strokeColor": "#0b7285", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 25, + "version": 1, + "versionNonce": 25, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "t6", + "type": "text", + "x": 103, + "y": 597, + "width": 634, + "height": 72, + "text": "Scheduler loop\n\nUTC daily tick -> enqueue digest_daily(yesterday) + flush_stale(today)\nflush_stale scans stale buffers and enqueues force seal jobs\nworkers consume these through the same mem_tree_jobs pipeline", + "fontSize": 18, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 68, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 26, + "version": 1, + "versionNonce": 26, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s1", + "type": "rectangle", + "x": 875, + "y": 585, + "width": 130, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#fff3bf", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 27, + "version": 1, + "versionNonce": 27, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st1", + "type": "text", + "x": 891, + "y": 607, + "width": 98, + "height": 24, + "text": "pending_extraction", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 28, + "version": 1, + "versionNonce": 28, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s2", + "type": "rectangle", + "x": 1045, + "y": 585, + "width": 110, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#d3f9d8", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 29, + "version": 1, + "versionNonce": 29, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st2", + "type": "text", + "x": 1069, + "y": 607, + "width": 62, + "height": 24, + "text": "admitted", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 30, + "version": 1, + "versionNonce": 30, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s3", + "type": "rectangle", + "x": 1195, + "y": 585, + "width": 110, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#d0ebff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 31, + "version": 1, + "versionNonce": 31, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st3", + "type": "text", + "x": 1223, + "y": 607, + "width": 54, + "height": 24, + "text": "buffered", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 32, + "version": 1, + "versionNonce": 32, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s4", + "type": "rectangle", + "x": 1345, + "y": 585, + "width": 110, + "height": 70, + "strokeColor": "#495057", + "backgroundColor": "#e5dbff", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 33, + "version": 1, + "versionNonce": 33, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st4", + "type": "text", + "x": 1375, + "y": 607, + "width": 50, + "height": 24, + "text": "sealed", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 20, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 34, + "version": 1, + "versionNonce": 34, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "s5", + "type": "rectangle", + "x": 1045, + "y": 665, + "width": 110, + "height": 36, + "strokeColor": "#c92a2a", + "backgroundColor": "#ffe3e3", + "fillStyle": "solid", + "strokeWidth": 2, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "roundness": { "type": 3 }, + "seed": 35, + "version": 1, + "versionNonce": 35, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "st5", + "type": "text", + "x": 1074, + "y": 672, + "width": 52, + "height": 20, + "text": "dropped", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "center", + "verticalAlign": "top", + "baseline": 16, + "strokeColor": "#1e1e1e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 36, + "version": 1, + "versionNonce": 36, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "life-note", + "type": "text", + "x": 1185, + "y": 665, + "width": 390, + "height": 36, + "text": "extract_chunk decides admitted vs dropped. append_buffer moves admitted leaves into buffers. seal creates summaries and parent links.", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 32, + "strokeColor": "#495057", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 37, + "version": 1, + "versionNonce": 37, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + }, + { + "id": "a1", + "type": "arrow", + "x": 315, + "y": 305, + "width": 115, + "height": 0, + "points": [[0, 0], [115, 0]], + "strokeColor": "#2f9e44", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 38, + "version": 1, + "versionNonce": 38, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [115, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a2", + "type": "arrow", + "x": 685, + "y": 305, + "width": 125, + "height": 0, + "points": [[0, 0], [125, 0]], + "strokeColor": "#e67700", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 39, + "version": 1, + "versionNonce": 39, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [125, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a3", + "type": "arrow", + "x": 1115, + "y": 305, + "width": 125, + "height": 0, + "points": [[0, 0], [125, 0]], + "strokeColor": "#9c36b5", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 40, + "version": 1, + "versionNonce": 40, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [125, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a4", + "type": "arrow", + "x": 430, + "y": 575, + "width": 70, + "height": 120, + "points": [[0, 0], [70, -120]], + "strokeColor": "#0b7285", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 41, + "version": 1, + "versionNonce": 41, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [70, -120], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a5", + "type": "arrow", + "x": 1005, + "y": 620, + "width": 40, + "height": 0, + "points": [[0, 0], [40, 0]], + "strokeColor": "#2b8a3e", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 42, + "version": 1, + "versionNonce": 42, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [40, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a6", + "type": "arrow", + "x": 1155, + "y": 620, + "width": 40, + "height": 0, + "points": [[0, 0], [40, 0]], + "strokeColor": "#1971c2", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 43, + "version": 1, + "versionNonce": 43, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [40, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a7", + "type": "arrow", + "x": 1305, + "y": 620, + "width": 40, + "height": 0, + "points": [[0, 0], [40, 0]], + "strokeColor": "#9c36b5", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 44, + "version": 1, + "versionNonce": 44, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [40, 0], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "a8", + "type": "arrow", + "x": 1100, + "y": 655, + "width": 0, + "height": 10, + "points": [[0, 0], [0, 10]], + "strokeColor": "#c92a2a", + "backgroundColor": "transparent", + "fillStyle": "solid", + "strokeWidth": 3, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 45, + "version": 1, + "versionNonce": 45, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false, + "lastCommittedPoint": [0, 10], + "startBinding": null, + "endBinding": null, + "startArrowhead": null, + "endArrowhead": "arrow" + }, + { + "id": "foot", + "type": "text", + "x": 40, + "y": 750, + "width": 1540, + "height": 60, + "text": "Job kinds in play: extract_chunk -> append_buffer(source) -> optional seal -> topic_route -> optional append_buffer(topic). Independent background flow: scheduler -> digest_daily / flush_stale -> seal. All paths go through mem_tree_jobs, so retries, dedupe, stale lock recovery, and worker wakeups stay centralized.", + "fontSize": 16, + "fontFamily": 3, + "textAlign": "left", + "verticalAlign": "top", + "baseline": 56, + "strokeColor": "#343a40", + "backgroundColor": "#ffffff", + "fillStyle": "solid", + "strokeWidth": 1, + "strokeStyle": "solid", + "roughness": 0, + "opacity": 100, + "angle": 0, + "seed": 46, + "version": 1, + "versionNonce": 46, + "isDeleted": false, + "groupIds": [], + "boundElements": null, + "updated": 1, + "link": null, + "locked": false + } + ], + "appState": { + "gridSize": null, + "viewBackgroundColor": "#ffffff" + }, + "files": {} +} diff --git a/src/bin/slack_backfill.rs b/src/bin/slack_backfill.rs index fffeffd7871..d6529f6e811 100644 --- a/src/bin/slack_backfill.rs +++ b/src/bin/slack_backfill.rs @@ -96,6 +96,21 @@ struct Cli { /// re-attempts cascade on the next append. #[arg(long = "seal-probe", default_value_t = false)] seal_probe: bool, + + /// Fire N back-to-back `SLACK_FETCH_CONVERSATION_HISTORY` calls + /// against the first listed channel and report a per-call tally + /// of {success, ratelimit, other-failure} + total duration. No + /// pacing by default (see --probe-pacing-ms), no ingestion. Used + /// to characterise Composio/Slack quota behaviour without + /// touching the memory tree. + #[arg(long = "probe-ratelimit")] + probe_ratelimit: Option, + + /// Sleep this many milliseconds between probe calls. 0 = fire + /// back-to-back (default). Use to find the threshold at which + /// rate-limits stop firing. + #[arg(long = "probe-pacing-ms", default_value_t = 0)] + probe_pacing_ms: u64, } #[tokio::main] @@ -243,6 +258,130 @@ async fn main() -> Result<()> { return Ok(()); } + if let Some(n) = cli.probe_ratelimit { + // Pure quota probe: fire N back-to-back + // SLACK_FETCH_CONVERSATION_HISTORY calls against the first + // discoverable channel. No pacing, no retry, no ingest. Reports + // a per-call status table + summary so we can characterise + // Composio/Slack rate-limit behaviour without contaminating the + // memory tree or burning extra quota on retries. + log::info!("[probe-ratelimit] requesting one channel via SLACK_LIST_CONVERSATIONS"); + let list_resp = client + .execute_tool( + "SLACK_LIST_CONVERSATIONS", + Some(serde_json::json!({ "exclude_archived": true, "limit": 1 })), + ) + .await + .map_err(|e| anyhow::anyhow!("SLACK_LIST_CONVERSATIONS failed: {e:#}"))?; + if !list_resp.successful { + anyhow::bail!( + "SLACK_LIST_CONVERSATIONS returned non-success: {:?}", + list_resp.error + ); + } + let channel_id = ["/data/channels/0/id", "/channels/0/id", "/data/0/id"] + .iter() + .find_map(|p| list_resp.data.pointer(p).and_then(|v| v.as_str())) + .map(str::to_string) + .ok_or_else(|| { + anyhow::anyhow!( + "could not find a channel id in SLACK_LIST_CONVERSATIONS response: {}", + serde_json::to_string(&list_resp.data).unwrap_or_default() + ) + })?; + log::info!("[probe-ratelimit] firing {n} calls against channel={channel_id}"); + + #[derive(Debug)] + enum Outcome { + Ok, + Ratelimit, + OtherFail(String), + Transport(String), + } + let mut outcomes: Vec<(u32, std::time::Duration, Outcome)> = Vec::with_capacity(n as usize); + let probe_started = Instant::now(); + for i in 1..=n { + if i > 1 && cli.probe_pacing_ms > 0 { + tokio::time::sleep(std::time::Duration::from_millis(cli.probe_pacing_ms)).await; + } + let t0 = Instant::now(); + let resp = client + .execute_tool( + "SLACK_FETCH_CONVERSATION_HISTORY", + Some(serde_json::json!({ "channel": channel_id, "limit": 1000 })), + ) + .await; + let dt = t0.elapsed(); + let outcome = match resp { + Err(e) => Outcome::Transport(format!("{e:#}")), + Ok(r) if r.successful => Outcome::Ok, + Ok(r) => { + let err = r.error.as_deref().unwrap_or("provider failure"); + if err.contains("ratelimited") + || err.contains("rate_limit") + || err.contains("rate limit") + { + log::warn!( + "[probe-ratelimit] call {i} ratelimited; body: {}", + serde_json::to_string(&r.data).unwrap_or_default() + ); + Outcome::Ratelimit + } else { + Outcome::OtherFail(err.to_string()) + } + } + }; + log::info!( + "[probe-ratelimit] call {i}/{n} took {:.2}s -> {:?}", + dt.as_secs_f64(), + outcome + ); + outcomes.push((i, dt, outcome)); + } + let total = probe_started.elapsed(); + + let ok = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::Ok)) + .count(); + let rl = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::Ratelimit)) + .count(); + let other = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::OtherFail(_))) + .count(); + let transport = outcomes + .iter() + .filter(|(_, _, o)| matches!(o, Outcome::Transport(_))) + .count(); + let avg_ms = if !outcomes.is_empty() { + outcomes.iter().map(|(_, d, _)| d.as_millis()).sum::() / outcomes.len() as u128 + } else { + 0 + }; + + println!("=== probe-ratelimit summary ==="); + println!("channel: {channel_id}"); + println!("calls fired: {n}"); + println!("total duration: {:.2}s", total.as_secs_f64()); + println!("avg per call: {avg_ms} ms"); + println!("successful: {ok}"); + println!("ratelimited: {rl}"); + println!("other failures: {other}"); + println!("transport errors: {transport}"); + if rl > 0 { + let first_rl = outcomes + .iter() + .find(|(_, _, o)| matches!(o, Outcome::Ratelimit)) + .map(|(i, _, _)| *i) + .unwrap_or(0); + println!("first ratelimit: call #{first_rl}"); + } + return Ok(()); + } + let connections = client .list_connections() .await diff --git a/src/core/jsonrpc.rs b/src/core/jsonrpc.rs index 90d5712b680..21fcb45cfd4 100644 --- a/src/core/jsonrpc.rs +++ b/src/core/jsonrpc.rs @@ -805,7 +805,10 @@ async fn run_server_inner( /// /// Guarded by `std::sync::Once` so repeated calls to `bootstrap_skill_runtime` /// are safe and idempotent. -fn register_domain_subscribers(workspace_dir: std::path::PathBuf) { +fn register_domain_subscribers( + workspace_dir: std::path::PathBuf, + config: crate::openhuman::config::Config, +) { use std::sync::{Arc, Once}; static REGISTERED: Once = Once::new(); @@ -840,6 +843,7 @@ fn register_domain_subscribers(workspace_dir: std::path::PathBuf) { } crate::openhuman::composio::register_composio_trigger_subscriber(); crate::openhuman::composio::start_periodic_sync(); + crate::openhuman::memory::tree::jobs::start(config.clone()); // Restart requests go through a subscriber so every trigger path shares // the same respawn logic. @@ -880,7 +884,7 @@ pub async fn bootstrap_skill_runtime() { // Register domain subscribers for cross-module event handling. // Uses a Once guard so repeated calls to bootstrap_skill_runtime() // cannot double-subscribe. - register_domain_subscribers(workspace_dir.clone()); + register_domain_subscribers(workspace_dir.clone(), cfg.clone()); // --- Sub-agent definition registry bootstrap --- // Loads built-in archetype definitions plus any custom TOML files diff --git a/src/openhuman/composio/providers/gmail/ingest.rs b/src/openhuman/composio/providers/gmail/ingest.rs new file mode 100644 index 00000000000..95a0d1d0be1 --- /dev/null +++ b/src/openhuman/composio/providers/gmail/ingest.rs @@ -0,0 +1,443 @@ +//! Gmail → memory tree ingest plumbing. +//! +//! Owns the conversion from a page of `GMAIL_FETCH_EMAILS` slim-envelope +//! messages (post-processed by [`super::post_process`]) into +//! [`EmailThread`] batches grouped by `(sender, threadId)`, then drives +//! [`memory::tree::ingest::ingest_email`] per thread with +//! [`GmailMarkdownStyle::Standard`]. +//! +//! Mirrors the bin's `bucket_by_sender_and_thread` grouping and the +//! per-thread rendering tested in `gmail-fetch-emails.rs`. Source-id is +//! per-inbox (`gmail:{connection_id}`), so every thread from one +//! connection lands in the same source tree. +//! +//! Idempotency: chunk IDs are content-hashed inside the memory tree, so +//! re-ingesting a previously-seen Gmail message is an UPSERT — buffer +//! token_sum may drift if content changes (rare for sealed mail), but +//! the tree's seal cascade handles that on next append. + +use std::collections::BTreeMap; + +use anyhow::Result; +use serde_json::Value; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::canonicalize::email::{ + EmailMessage, EmailThread, GmailMarkdownStyle, +}; +use crate::openhuman::memory::tree::canonicalize::email_clean::{ + extract_email, parse_message_date, +}; +use crate::openhuman::memory::tree::ingest::{ingest_email, IngestResult}; + +/// Provider name embedded in the canonical email-thread header. Matches +/// the value `memory::tree::retrieval::source::PLATFORM_KINDS` expects. +pub const GMAIL_PROVIDER: &str = "gmail"; + +/// Tags attached to every Gmail-ingested chunk. Stable list — retrieval +/// callers filter on these. +pub const DEFAULT_TAGS: &[&str] = &["gmail", "ingested"]; + +/// Inner map: `threadId → ordered messages`. +type ThreadBucket<'a> = BTreeMap>; + +/// Group raw page messages by `(sender_email, thread_id)`. Within a +/// thread, messages sort ascending by date so each rendered thread reads +/// chronologically. Mirrors the same-named function in +/// `bin/gmail_fetch_emails.rs` — kept aligned so bin output and +/// production canonicalisation see the same bucketing. +pub(crate) fn bucket_by_sender_and_thread( + msgs: &[Value], +) -> BTreeMap> { + let mut out: BTreeMap> = BTreeMap::new(); + for m in msgs { + let from = m.get("from").and_then(|v| v.as_str()).unwrap_or(""); + let sender = extract_email(from) + .map(|s| s.to_lowercase()) + .unwrap_or_else(|| { + if from.trim().is_empty() { + "unknown".to_string() + } else { + from.trim().to_lowercase() + } + }); + let thread = m + .get("threadId") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(str::to_string) + .or_else(|| { + m.get("id") + .and_then(|v| v.as_str()) + .map(|s| format!("solo:{s}")) + }) + .unwrap_or_else(|| "unknown".to_string()); + out.entry(sender) + .or_default() + .entry(thread) + .or_default() + .push(m); + } + for threads in out.values_mut() { + for msgs in threads.values_mut() { + msgs.sort_by_key(|m| { + parse_message_date(m).map(|d| d.timestamp()).unwrap_or(0) + }); + } + } + out +} + +/// Build an [`EmailMessage`] from a raw slim-envelope JSON message. +/// Returns `None` when the message has no parseable date — the rest of +/// the pipeline can't sort or canonicalise without one. +pub(crate) fn raw_to_email_message(raw: &Value) -> Option { + let id = raw + .get("id") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .unwrap_or(""); + let from = raw + .get("from") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let to = parse_address_list(raw.get("to")); + let cc = parse_address_list(raw.get("cc")); + let subject = raw + .get("subject") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let sent_at = parse_message_date(raw)?; + let body = raw + .get("markdown") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let source_ref = if id.is_empty() { + None + } else { + Some(format!("gmail://msg/{id}")) + }; + Some(EmailMessage { + from, + to, + cc, + subject, + sent_at, + body, + source_ref, + }) +} + +/// Parse the `to` / `cc` field which Composio surfaces as either a +/// JSON array of strings or a single comma-separated string. Empty +/// entries are dropped. +fn parse_address_list(v: Option<&Value>) -> Vec { + match v { + Some(Value::Array(arr)) => arr + .iter() + .filter_map(|s| s.as_str()) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(), + Some(Value::String(s)) => s + .split(',') + .map(|p| p.trim().to_string()) + .filter(|p| !p.is_empty()) + .collect(), + _ => Vec::new(), + } +} + +/// Ingest a page of raw Gmail messages into the memory tree under one +/// inbox. Each `(sender, thread)` group becomes one [`EmailThread`] with +/// [`GmailMarkdownStyle::Standard`], handed to +/// [`ingest_email`] which fans out to the chunker + scorer + source +/// tree downstream. +/// +/// Returns the total number of chunks written across all threads so +/// callers can surface counts in logs / outcomes. Per-thread errors are +/// logged and swallowed — one bad thread should not abort the whole +/// page (the next sync re-fetches it via the date-cursor). +pub async fn ingest_page_into_memory_tree( + config: &Config, + source_id: &str, + owner: &str, + page_messages: &[Value], +) -> Result { + if page_messages.is_empty() { + return Ok(0); + } + let buckets = bucket_by_sender_and_thread(page_messages); + let mut total_chunks = 0usize; + let mut total_threads = 0usize; + for (sender, threads) in &buckets { + for (thread_id, raw_msgs) in threads { + let messages: Vec = + raw_msgs.iter().filter_map(|raw| raw_to_email_message(raw)).collect(); + if messages.is_empty() { + log::debug!( + "[composio:gmail][ingest] skipping empty thread sender={sender} thread={thread_id}" + ); + continue; + } + let thread_subject = pick_thread_subject(&messages); + log::info!( + "[composio:gmail][ingest] thread sender={} thread_id={} messages={} source_id={}", + sender, + thread_id, + messages.len(), + source_id + ); + let thread = EmailThread { + provider: GMAIL_PROVIDER.to_string(), + thread_subject, + messages, + gmail_style: Some(GmailMarkdownStyle::Standard), + }; + let tags = DEFAULT_TAGS.iter().map(|s| (*s).to_string()).collect(); + match ingest_email(config, source_id, owner, tags, thread).await { + Ok(IngestResult { chunks_written, .. }) => { + total_chunks += chunks_written; + total_threads += 1; + } + Err(e) => { + log::warn!( + "[composio:gmail][ingest] ingest_email failed sender={} thread={} err={:#}", + sender, + thread_id, + e + ); + } + } + } + } + log::info!( + "[composio:gmail][ingest] page_done source_id={source_id} threads={total_threads} chunks={total_chunks}" + ); + Ok(total_chunks) +} + +/// Strip "Re:" / "Fwd:" prefixes from the head message's subject so +/// every message in a thread shares one canonical thread subject. Falls +/// back to "(no subject)" when empty. +fn pick_thread_subject(messages: &[EmailMessage]) -> String { + let raw = messages + .first() + .map(|m| m.subject.trim().to_string()) + .unwrap_or_default(); + let stripped = strip_reply_prefixes(&raw); + if stripped.is_empty() { + "(no subject)".to_string() + } else { + stripped + } +} + +/// Iteratively strip `Re:` / `Fwd:` / `Fw:` prefixes (case-insensitive, +/// optional whitespace) from the front of a subject. Stops once a pass +/// removes nothing. +fn strip_reply_prefixes(subject: &str) -> String { + let mut s = subject.trim().to_string(); + loop { + let lower = s.to_ascii_lowercase(); + let stripped = if lower.starts_with("re:") { + Some(&s[3..]) + } else if lower.starts_with("fwd:") { + Some(&s[4..]) + } else if lower.starts_with("fw:") { + Some(&s[3..]) + } else { + None + }; + match stripped { + Some(rest) => { + let trimmed = rest.trim_start().to_string(); + if trimmed == s { + return s; + } + s = trimmed; + } + None => return s, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn bucket_groups_by_sender_and_thread() { + let msgs = vec![ + json!({ + "id": "m1", "threadId": "t1", + "from": "Alice ", + "subject": "Hi", + "date": "2026-04-21T10:00:00Z", + "markdown": "first", + }), + json!({ + "id": "m2", "threadId": "t1", + "from": "Alice ", + "subject": "Re: Hi", + "date": "2026-04-21T11:00:00Z", + "markdown": "second", + }), + json!({ + "id": "m3", "threadId": "t2", + "from": "bob@example.com", + "subject": "Other", + "date": "2026-04-22T09:00:00Z", + "markdown": "third", + }), + ]; + let buckets = bucket_by_sender_and_thread(&msgs); + assert_eq!(buckets.len(), 2); + assert_eq!(buckets["alice@example.com"]["t1"].len(), 2); + assert_eq!(buckets["bob@example.com"]["t2"].len(), 1); + // Within a thread, messages sort ascending by date. + let t1 = &buckets["alice@example.com"]["t1"]; + assert_eq!(t1[0].get("id").unwrap().as_str().unwrap(), "m1"); + assert_eq!(t1[1].get("id").unwrap().as_str().unwrap(), "m2"); + } + + #[test] + fn solo_thread_id_when_threadid_missing() { + let msgs = vec![json!({ + "id": "abc", + "from": "noreply@github.com", + "subject": "x", + "date": "2026-04-21T10:00:00Z", + "markdown": "body", + })]; + let buckets = bucket_by_sender_and_thread(&msgs); + assert!(buckets["noreply@github.com"].contains_key("solo:abc")); + } + + #[test] + fn unknown_sender_fallback_when_from_blank() { + let msgs = vec![json!({ + "id": "abc", "threadId": "t1", + "from": "", + "subject": "x", + "date": "2026-04-21T10:00:00Z", + "markdown": "body", + })]; + let buckets = bucket_by_sender_and_thread(&msgs); + assert!(buckets.contains_key("unknown")); + } + + #[test] + fn raw_to_email_message_parses_slim_envelope() { + let raw = json!({ + "id": "m1", + "from": "Alice ", + "to": "me@example.com", + "cc": "team@example.com", + "subject": "Phoenix kickoff", + "date": "2026-04-21T10:00:00Z", + "markdown": "Let's ship Phoenix.", + }); + let msg = raw_to_email_message(&raw).unwrap(); + assert_eq!(msg.from, "Alice "); + assert_eq!(msg.to, vec!["me@example.com"]); + assert_eq!(msg.cc, vec!["team@example.com"]); + assert_eq!(msg.subject, "Phoenix kickoff"); + assert_eq!(msg.body, "Let's ship Phoenix."); + assert_eq!(msg.source_ref.as_deref(), Some("gmail://msg/m1")); + } + + #[test] + fn raw_to_email_message_handles_to_array() { + let raw = json!({ + "id": "m1", + "from": "a@x", + "to": ["b@x", "c@x"], + "subject": "x", + "date": "2026-04-21T10:00:00Z", + "markdown": "body", + }); + let msg = raw_to_email_message(&raw).unwrap(); + assert_eq!(msg.to, vec!["b@x", "c@x"]); + } + + #[test] + fn raw_to_email_message_handles_comma_separated_to_string() { + let raw = json!({ + "id": "m1", + "from": "a@x", + "to": "b@x, c@x ,d@x", + "subject": "x", + "date": "2026-04-21T10:00:00Z", + "markdown": "body", + }); + let msg = raw_to_email_message(&raw).unwrap(); + assert_eq!(msg.to, vec!["b@x", "c@x", "d@x"]); + } + + #[test] + fn raw_to_email_message_returns_none_on_unparseable_date() { + let raw = json!({ + "id": "m1", + "from": "a@x", + "subject": "x", + "date": "not-a-date", + "markdown": "body", + }); + assert!(raw_to_email_message(&raw).is_none()); + } + + #[test] + fn raw_to_email_message_drops_source_ref_when_id_empty() { + let raw = json!({ + "id": "", + "from": "a@x", + "subject": "x", + "date": "2026-04-21T10:00:00Z", + "markdown": "body", + }); + let msg = raw_to_email_message(&raw).unwrap(); + assert!(msg.source_ref.is_none()); + } + + #[test] + fn strip_reply_prefixes_removes_iterated() { + assert_eq!(strip_reply_prefixes("Re: Re: Hi"), "Hi"); + assert_eq!(strip_reply_prefixes("Fwd: Re: Status"), "Status"); + assert_eq!(strip_reply_prefixes("RE: Question"), "Question"); + assert_eq!(strip_reply_prefixes("Fw: alert"), "alert"); + assert_eq!(strip_reply_prefixes("Plain subject"), "Plain subject"); + } + + #[test] + fn pick_thread_subject_strips_reply_prefixes() { + let messages = vec![EmailMessage { + from: "a@x".into(), + to: vec![], + cc: vec![], + subject: "Re: Re: Phoenix kickoff".into(), + sent_at: chrono::Utc::now(), + body: "body".into(), + source_ref: None, + }]; + assert_eq!(pick_thread_subject(&messages), "Phoenix kickoff"); + } + + #[test] + fn pick_thread_subject_falls_back_to_no_subject() { + let messages = vec![EmailMessage { + from: "a@x".into(), + to: vec![], + cc: vec![], + subject: " ".into(), + sent_at: chrono::Utc::now(), + body: "body".into(), + source_ref: None, + }]; + assert_eq!(pick_thread_subject(&messages), "(no subject)"); + } +} diff --git a/src/openhuman/memory/tree/canonicalize/email_clean.rs b/src/openhuman/memory/tree/canonicalize/email_clean.rs new file mode 100644 index 00000000000..50d3d889344 --- /dev/null +++ b/src/openhuman/memory/tree/canonicalize/email_clean.rs @@ -0,0 +1,380 @@ +//! Shared email rendering + cleaning helpers. +//! +//! Used by both [`canonicalize::email`](super::email) (when rendering the +//! `GmailMarkdownStyle::Standard` shape) and the `gmail-fetch-emails` bin +//! (which writes per-sender markdown digests to disk). Lifted out of the +//! bin so the bin's behaviour and the production canonicaliser stay +//! byte-identical on body cleanup. +//! +//! The module is intentionally pure-string-oriented + a single +//! `serde_json::Value` helper (`parse_message_date`) used by callers that +//! work directly off Gmail's slim envelope JSON. Nothing here depends on +//! the memory-tree types — that keeps the helpers reusable. + +use chrono::{DateTime, NaiveDate, Utc}; +use serde_json::Value; + +/// Two-stage cleanup applied to each message body before it gets +/// blockquoted into a digest: +/// +/// 1. **Drop quoted reply chains** — once a message contains a +/// `On , wrote:` preamble, an `Original Message` / +/// `Forwarded message` separator, or a run of three+ consecutive +/// `>`-prefixed lines, everything from that point onward is the +/// parent message we already render directly above. +/// 2. **Drop footer noise** — `Unsubscribe`, `View in browser`, +/// copyright lines, legal disclaimers, and address blocks. We cut +/// at the first line containing any of [`FOOTER_TRIGGERS`]. +/// +/// The two passes run in order so a quoted-chain preamble below a +/// "view in browser" line still gets stripped on its own merits even +/// if the footer pass missed it. +pub fn clean_body(raw: &str) -> String { + let stage1 = drop_reply_chain(raw); + let stage2 = drop_footer_noise(&stage1); + collapse_blank_runs(stage2.trim()) +} + +/// Substrings that, when matched (case-insensitive) anywhere on a +/// line, mark the start of footer / boilerplate territory. Conservative +/// list — every entry should be unambiguous noise that wouldn't +/// reasonably appear inside real prose. +const FOOTER_TRIGGERS: &[&str] = &[ + "unsubscribe", + "view in browser", + "view this email in your browser", + "view it in your browser", + "update your email settings", + "manage your subscription", + "manage preferences", + "email preferences", + "you are receiving this email because", + "you received this email because", + "you're receiving this email because", + "to stop receiving", + "all rights reserved", + "© 20", + "(c) 20", + "copyright 20", + "powered by mailchimp", + "sent via sendgrid", + "this email and any files", + "confidentiality notice", + "if you are not the intended recipient", + "this communication may contain", +]; + +/// Strip quoted reply chains. See [`clean_body`] for details. +pub fn drop_reply_chain(s: &str) -> String { + let mut offset = 0usize; + let mut quoted_run_start: Option = None; + let mut quoted_run_len = 0u32; + + for line in s.split_inclusive('\n') { + let trimmed = line.trim(); + let lower = trimmed.to_ascii_lowercase(); + + // Explicit reply / forward markers. + let is_preamble = (lower.starts_with("on ") && lower.contains(" wrote:")) + || lower.contains("---------- forwarded message") + || lower.contains("----- original message") + || lower.contains("--------- original message") + || lower.contains("--- forwarded by"); + if is_preamble { + return s[..offset].trim_end().to_string(); + } + + // Three+ consecutive lines starting with `>` is a quoted + // reply chain in disguise (some clients de-quote on send). + // Treat the start of the run as the cut point. + if trimmed.starts_with('>') { + if quoted_run_start.is_none() { + quoted_run_start = Some(offset); + quoted_run_len = 1; + } else { + quoted_run_len += 1; + } + if quoted_run_len >= 3 { + let cut = quoted_run_start.unwrap_or(offset); + return s[..cut].trim_end().to_string(); + } + } else if !trimmed.is_empty() { + // Reset on a non-empty, non-quoted line. Blank lines + // don't break a quote run because senders often interleave + // them. + quoted_run_start = None; + quoted_run_len = 0; + } + + offset += line.len(); + } + s.to_string() +} + +/// Strip everything from the first line containing a footer trigger +/// onward. See [`FOOTER_TRIGGERS`] for the matched list. +pub fn drop_footer_noise(s: &str) -> String { + let mut offset = 0usize; + for line in s.split_inclusive('\n') { + let lower = line.to_ascii_lowercase(); + if FOOTER_TRIGGERS.iter().any(|t| lower.contains(t)) { + return s[..offset].trim_end().to_string(); + } + offset += line.len(); + } + s.to_string() +} + +/// Collapse runs of 2+ blank lines into a single blank line. Trims +/// trailing newlines. +pub fn collapse_blank_runs(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + let mut blank = 0u32; + for line in s.lines() { + if line.trim().is_empty() { + blank += 1; + if blank <= 1 { + out.push('\n'); + } + } else { + blank = 0; + out.push_str(line); + out.push('\n'); + } + } + while out.ends_with('\n') { + out.pop(); + } + out +} + +/// Truncate a body to at most `max_chars` characters, appending `…` when +/// the body is longer. Trims first so leading/trailing whitespace doesn't +/// count against the budget. +pub fn truncate_body(body: &str, max_chars: usize) -> String { + let trimmed = body.trim(); + if trimmed.chars().count() <= max_chars { + return trimmed.to_string(); + } + let mut out: String = trimmed.chars().take(max_chars).collect(); + out.push('…'); + out +} + +/// Escape only the few markdown chars that would visibly break the +/// header/inline contexts we use (#, |, *, _, `). Newlines collapse to +/// spaces. We leave most punctuation alone — the body is rendered as a +/// blockquote anyway. +pub fn md_escape(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + for ch in s.chars() { + match ch { + '\\' | '`' | '*' | '_' | '|' => { + out.push('\\'); + out.push(ch); + } + '\n' | '\r' => out.push(' '), + _ => out.push(ch), + } + } + out +} + +/// Pull the `` portion out of a `From` header, returning +/// just the bare email address. Falls back to `None` when no `<…>` +/// brackets exist; in that case the caller may use the raw From field. +pub fn extract_email(from: &str) -> Option { + let s = from.trim(); + if let (Some(start), Some(end)) = (s.rfind('<'), s.rfind('>')) { + if start < end { + let inner = s[start + 1..end].trim(); + if inner.contains('@') { + return Some(inner.to_string()); + } + } + } + if s.contains('@') && !s.contains(' ') { + return Some(s.to_string()); + } + None +} + +/// If `s` starts with a 3-letter day-of-week prefix (`Mon, `, `Tue, `, …), +/// return the remainder; otherwise `None`. Used to feed a strict-rfc2822 +/// reject into a lenient retry. +fn strip_day_of_week_prefix(s: &str) -> Option<&str> { + const DAYS: &[&str] = &["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]; + let (prefix, rest) = s.split_once(", ")?; + if DAYS.iter().any(|d| d.eq_ignore_ascii_case(prefix)) { + Some(rest) + } else { + None + } +} + +/// Try a sequence of common date formats. Composio's slim envelope sets +/// `date` from `messageTimestamp` (often ISO 8601 or epoch ms) when +/// present, falling back to the raw `Date:` header (RFC 2822). Operates +/// on the raw `serde_json::Value` so callers that work off the slim +/// envelope JSON don't have to reshape it first. +pub fn parse_message_date(m: &Value) -> Option> { + let raw = m.get("date")?; + if let Some(s) = raw.as_str() { + let s = s.trim(); + if s.is_empty() { + return None; + } + // Epoch millis as a string? + if let Ok(ms) = s.parse::() { + return DateTime::from_timestamp_millis(ms); + } + if let Ok(dt) = DateTime::parse_from_rfc3339(s) { + return Some(dt.with_timezone(&Utc)); + } + if let Ok(dt) = DateTime::parse_from_rfc2822(s) { + return Some(dt.with_timezone(&Utc)); + } + // Lenient RFC 2822 fallback: strict `parse_from_rfc2822` rejects + // mismatched day-of-week (e.g. `Mon, 21 Apr 2026 …` when Apr 21 + // 2026 is a Tuesday). Real-world MTAs occasionally send these + // with a wrong day-name. Strip a `, ` prefix and retry + // with the rfc2822 body format. Keeps the date if everything + // else is sane. + if let Some(rest) = strip_day_of_week_prefix(s) { + if let Ok(dt) = DateTime::parse_from_str(rest, "%d %b %Y %H:%M:%S %z") { + return Some(dt.with_timezone(&Utc)); + } + } + if let Ok(d) = NaiveDate::parse_from_str(s, "%Y-%m-%d") { + return d.and_hms_opt(0, 0, 0).map(|n| n.and_utc()); + } + } + if let Some(ms) = raw.as_i64() { + return DateTime::from_timestamp_millis(ms); + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn drop_reply_chain_strips_on_x_wrote_preamble() { + let body = "Sounds good — let's do Tuesday.\n\nOn Mon, Apr 22, 2026 at 10:00 AM, Alice wrote:\n> Tuesday or Wednesday?\n> Let me know."; + let cleaned = drop_reply_chain(body); + assert_eq!(cleaned.trim(), "Sounds good — let's do Tuesday."); + } + + #[test] + fn drop_reply_chain_strips_forwarded_separator() { + let body = "FYI.\n\n---------- Forwarded message ---------\nFrom: bob\nSubject: hi"; + assert_eq!(drop_reply_chain(body).trim(), "FYI."); + } + + #[test] + fn drop_reply_chain_strips_consecutive_quoted_run() { + let body = "Thanks for the update.\n\n> earlier line 1\n> earlier line 2\n> earlier line 3\n> earlier line 4"; + assert_eq!(drop_reply_chain(body).trim(), "Thanks for the update."); + } + + #[test] + fn drop_reply_chain_keeps_short_quote() { + // A single inline blockquote is fine — only 3+ consecutive lines trigger. + let body = "I think:\n> That sounds reasonable\n\nLet's proceed."; + let cleaned = drop_reply_chain(body); + assert!(cleaned.contains("Let's proceed")); + assert!(cleaned.contains("That sounds reasonable")); + } + + #[test] + fn drop_footer_noise_strips_unsubscribe_block() { + let body = "Big news: GPT-5.5 is here.\n\nRead more at example.com\n\nUnsubscribe | © 2026 OpenAI"; + let cleaned = drop_footer_noise(body); + assert!(cleaned.contains("GPT-5.5")); + assert!(!cleaned.to_ascii_lowercase().contains("unsubscribe")); + assert!(!cleaned.contains("©")); + } + + #[test] + fn drop_footer_noise_strips_legal_disclaimer() { + let body = "Action item — review by Friday.\n\nThis email and any files transmitted with it are confidential and intended solely for the use of the individual to whom they are addressed."; + let cleaned = drop_footer_noise(body); + assert_eq!(cleaned.trim(), "Action item — review by Friday."); + } + + #[test] + fn clean_body_combines_passes() { + let body = "Real content here.\n\nOn Mon, Apr 22, 2026, Alice wrote:\n> old stuff\n\nUnsubscribe"; + let cleaned = clean_body(body); + assert_eq!(cleaned, "Real content here."); + } + + #[test] + fn collapse_blank_runs_keeps_paragraph_breaks() { + let s = "a\n\n\n\nb\n\n\nc\n"; + assert_eq!(collapse_blank_runs(s), "a\n\nb\n\nc"); + } + + #[test] + fn truncate_body_adds_ellipsis() { + let s = "x".repeat(2000); + let t = truncate_body(&s, 1200); + assert!(t.ends_with('…')); + assert_eq!(t.chars().count(), 1201); + } + + #[test] + fn truncate_body_passthrough_when_short() { + let s = "hello"; + let t = truncate_body(s, 1200); + assert_eq!(t, "hello"); + } + + #[test] + fn md_escape_handles_special_chars() { + assert_eq!(md_escape("a*b_c"), "a\\*b\\_c"); + assert_eq!(md_escape("foo|bar"), "foo\\|bar"); + assert_eq!(md_escape("line1\nline2"), "line1 line2"); + assert_eq!(md_escape("plain text"), "plain text"); + } + + #[test] + fn extract_email_handles_both_forms() { + assert_eq!( + extract_email("Alice ").as_deref(), + Some("alice@example.com") + ); + assert_eq!( + extract_email("notify@github.com").as_deref(), + Some("notify@github.com") + ); + assert_eq!( + extract_email("\"Bot Name\" ").as_deref(), + Some("bot@x.io") + ); + assert!(extract_email("Alice").is_none()); + } + + #[test] + fn parse_message_date_handles_iso_and_rfc2822() { + let iso = json!({"date": "2026-04-21T10:00:00Z"}); + let rfc = json!({"date": "Mon, 21 Apr 2026 10:00:00 +0000"}); + let ms = json!({"date": 1745236800000_i64}); + let ms_str = json!({"date": "1745236800000"}); + let date_only = json!({"date": "2026-04-21"}); + assert!(parse_message_date(&iso).is_some()); + assert!(parse_message_date(&rfc).is_some()); + assert!(parse_message_date(&ms).is_some()); + assert!(parse_message_date(&ms_str).is_some()); + assert!(parse_message_date(&date_only).is_some()); + } + + #[test] + fn parse_message_date_returns_none_when_missing_or_blank() { + assert!(parse_message_date(&json!({})).is_none()); + assert!(parse_message_date(&json!({"date": ""})).is_none()); + assert!(parse_message_date(&json!({"date": " "})).is_none()); + } +} diff --git a/src/openhuman/memory/tree/global_tree/digest.rs b/src/openhuman/memory/tree/global_tree/digest.rs index e55187ee4a5..efabb0f0055 100644 --- a/src/openhuman/memory/tree/global_tree/digest.rs +++ b/src/openhuman/memory/tree/global_tree/digest.rs @@ -19,6 +19,8 @@ //! - Idempotency: if an L0 daily node already exists for the target day, //! return `DigestOutcome::Skipped` rather than emitting a duplicate. +use std::collections::BTreeSet; + use anyhow::{Context, Result}; use chrono::{DateTime, Duration, NaiveDate, TimeZone, Utc}; use rusqlite::OptionalExtension; @@ -154,6 +156,25 @@ pub async fn end_of_day_digest( .await .context("embed daily summary during end_of_day_digest")?; + // L0 daily node inherits entities/topics by union of contributing + // source-tree summaries. Each input was already labeled at source-tree + // seal time, so emergent themes don't need another extractor pass + // here — global is a sink; union preserves "days that mentioned X" + // retrieval without an extra LLM call. See LabelStrategy in + // source_tree::bucket_seal for the full design. + let mut entities_set: BTreeSet = BTreeSet::new(); + let mut topics_set: BTreeSet = BTreeSet::new(); + for inp in &inputs { + for e in &inp.entities { + entities_set.insert(e.clone()); + } + for t in &inp.topics { + topics_set.insert(t.clone()); + } + } + let daily_entities: Vec = entities_set.into_iter().collect(); + let daily_topics: Vec = topics_set.into_iter().collect(); + let now = Utc::now(); let daily_id = new_summary_id(0); let daily = SummaryNode { @@ -165,8 +186,8 @@ pub async fn end_of_day_digest( child_ids: inputs.iter().map(|i| i.id.clone()).collect(), content: output.content, token_count: output.token_count, - entities: output.entities, - topics: output.topics, + entities: daily_entities, + topics: daily_topics, time_range_start: day_start, time_range_end: day_end, score, diff --git a/src/openhuman/memory/tree/global_tree/digest_tests.rs b/src/openhuman/memory/tree/global_tree/digest_tests.rs index 0db1ca75002..109f6020cea 100644 --- a/src/openhuman/memory/tree/global_tree/digest_tests.rs +++ b/src/openhuman/memory/tree/global_tree/digest_tests.rs @@ -1,5 +1,7 @@ use super::*; -use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, +}; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; use crate::openhuman::memory::tree::source_tree::types::TreeStatus; @@ -76,8 +78,12 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime topics: vec![], score: 0.5, }; - append_leaf(cfg, &tree, &leaf1, &summariser).await.unwrap(); - append_leaf(cfg, &tree, &leaf2, &summariser).await.unwrap(); + append_leaf(cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); + append_leaf(cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); // 12k tokens > 10k budget → one L1 summary covering `ts`. } @@ -216,3 +222,155 @@ async fn seven_days_cascade_to_weekly_seal() { assert_eq!(t.max_level, 1); assert_eq!(t.status, TreeStatus::Active); } + +/// Seed a source tree whose sealed L1 summary carries the given entities +/// and topics. Entities are written into `mem_tree_entity_index` (where +/// seal-time hydration reads them); topics are stored on chunk metadata +/// tags. The seal then unions both into the L1 summary. +async fn seed_source_tree_with_labeled_l1( + cfg: &Config, + scope: &str, + ts: DateTime, + entities: Vec, + topics: Vec, +) { + use crate::openhuman::memory::tree::score::extract::EntityKind; + use crate::openhuman::memory::tree::score::resolver::CanonicalEntity; + use crate::openhuman::memory::tree::score::store::index_entity; + + let tree = get_or_create_source_tree(cfg, scope).unwrap(); + let summariser = InertSummariser::new(); + + let mut chunks: Vec = Vec::new(); + for seq in 0..2u32 { + chunks.push(Chunk { + id: chunk_id(SourceKind::Chat, scope, seq, "labeled-test"), + content: format!("labeled chunk {seq} in {scope}"), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: scope.into(), + owner: "alice".into(), + timestamp: ts, + time_range: (ts, ts), + tags: topics.clone(), + source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))), + }, + token_count: 6_000, + seq_in_source: seq, + created_at: ts, + }); + } + upsert_chunks(cfg, &chunks).unwrap(); + + for chunk in &chunks { + for entity_id in &entities { + let kind = entity_id + .split_once(':') + .map_or(EntityKind::Misc, |(k, _)| { + EntityKind::parse(k).unwrap_or(EntityKind::Misc) + }); + let surface = entity_id + .split_once(':') + .map_or(entity_id.as_str(), |(_, v)| v); + let e = CanonicalEntity { + canonical_id: entity_id.clone(), + kind, + surface: surface.to_string(), + span_start: 0, + span_end: surface.len() as u32, + score: 1.0, + }; + index_entity( + cfg, + &e, + &chunk.id, + "leaf", + ts.timestamp_millis(), + Some(scope), + ) + .unwrap(); + } + } + + // Two 6k-token leaves total 12k → exceeds L0 budget → seal fires on + // the second append, producing one L1 summary that unions all leaf + // labels (every leaf has the same set, so dedup yields the input set). + for chunk in &chunks { + let leaf = LeafRef { + chunk_id: chunk.id.clone(), + token_count: 6_000, + timestamp: ts, + content: chunk.content.clone(), + entities: entities.clone(), + topics: topics.clone(), + score: 0.5, + }; + append_leaf( + cfg, + &tree, + &leaf, + &summariser, + &LabelStrategy::UnionFromChildren, + ) + .await + .unwrap(); + } +} + +#[tokio::test] +async fn daily_digest_unions_labels_from_source_summaries() { + let (_tmp, cfg) = test_config(); + let summariser = InertSummariser::new(); + + let day = NaiveDate::from_ymd_opt(2025, 5, 1).unwrap(); + let ts = day.and_hms_opt(10, 0, 0).unwrap().and_utc(); + + // Source A's L1 carries (alice, phoenix-migration). Source B's L1 + // carries (bob, phoenix-migration, qa). The daily L0 should union to + // (alice, bob, phoenix-migration) for entities and (phoenix-migration, + // qa) for topics — overlap dedup'd. + seed_source_tree_with_labeled_l1( + &cfg, + "slack:#a", + ts, + vec!["email:alice@example.com".into(), "topic:phoenix".into()], + vec!["phoenix-migration".into()], + ) + .await; + seed_source_tree_with_labeled_l1( + &cfg, + "slack:#b", + ts, + vec!["person:bob".into(), "topic:phoenix".into()], + vec!["phoenix-migration".into(), "qa".into()], + ) + .await; + + let outcome = end_of_day_digest(&cfg, day, &summariser).await.unwrap(); + let daily_id = match outcome { + DigestOutcome::Emitted { daily_id, .. } => daily_id, + other => panic!("expected Emitted, got {other:?}"), + }; + + let daily = store::get_summary(&cfg, &daily_id).unwrap().unwrap(); + let entities: std::collections::BTreeSet<&str> = + daily.entities.iter().map(String::as_str).collect(); + let topics: std::collections::BTreeSet<&str> = + daily.topics.iter().map(String::as_str).collect(); + + assert!(entities.contains("email:alice@example.com")); + assert!(entities.contains("person:bob")); + assert!(entities.contains("topic:phoenix")); + assert_eq!( + entities.len(), + 3, + "expected 3 unique entities (deduped); got {entities:?}" + ); + assert!(topics.contains("phoenix-migration")); + assert!(topics.contains("qa")); + assert_eq!( + topics.len(), + 2, + "expected 2 unique topics (deduped); got {topics:?}" + ); +} diff --git a/src/openhuman/memory/tree/global_tree/recap.rs b/src/openhuman/memory/tree/global_tree/recap.rs index ddbccfde7c6..f92ae45d00f 100644 --- a/src/openhuman/memory/tree/global_tree/recap.rs +++ b/src/openhuman/memory/tree/global_tree/recap.rs @@ -159,7 +159,9 @@ fn assemble_recap(covering: &[&SummaryNode], level: u32) -> RecapOutput { mod tests { use super::*; use crate::openhuman::memory::tree::global_tree::digest::{end_of_day_digest, DigestOutcome}; - use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, + }; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; use crate::openhuman::memory::tree::store::upsert_chunks; @@ -247,6 +249,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); @@ -263,6 +266,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); diff --git a/src/openhuman/memory/tree/global_tree/seal.rs b/src/openhuman/memory/tree/global_tree/seal.rs index efe77b65868..a108e092819 100644 --- a/src/openhuman/memory/tree/global_tree/seal.rs +++ b/src/openhuman/memory/tree/global_tree/seal.rs @@ -11,6 +11,8 @@ //! `mem_tree_summaries` on both sides (children and output), since even L0 //! is a sealed summary node rather than a raw chunk. +use std::collections::BTreeSet; + use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; @@ -186,6 +188,26 @@ async fn seal_one_level( .await .context("summariser failed during global seal")?; + // Global-tree summaries inherit their entity/topic labels via union + // from their already-labeled inputs (source-tree summaries carry + // labels from the source-tree seal extractor; global L1+ inputs + // carry labels from this same union path one level down). We + // deliberately do NOT run an extractor on the daily/weekly/monthly + // synthesis: the inputs already cover what the summary represents, + // and global is a sink — no second-pass labeling earns its keep. + let mut entities_set: BTreeSet = BTreeSet::new(); + let mut topics_set: BTreeSet = BTreeSet::new(); + for inp in &inputs { + for e in &inp.entities { + entities_set.insert(e.clone()); + } + for t in &inp.topics { + topics_set.insert(t.clone()); + } + } + let node_entities: Vec = entities_set.into_iter().collect(); + let node_topics: Vec = topics_set.into_iter().collect(); + // Phase 4 (#710): embed BEFORE opening the write tx so an embedder // error aborts the cascade without half-committing the summary. let embedder = @@ -208,8 +230,8 @@ async fn seal_one_level( child_ids: buf.item_ids.clone(), content: output.content, token_count: output.token_count, - entities: output.entities, - topics: output.topics, + entities: node_entities, + topics: node_topics, time_range_start, time_range_end, score, diff --git a/src/openhuman/memory/tree/ingest.rs b/src/openhuman/memory/tree/ingest.rs index 91386ee3d19..c4dca65b0f5 100644 --- a/src/openhuman/memory/tree/ingest.rs +++ b/src/openhuman/memory/tree/ingest.rs @@ -1,14 +1,12 @@ -//! Ingest orchestrator (Phase 1 + Phase 2): +//! Ingest orchestrator for the async memory-tree pipeline. //! -//! canonicalise → chunk → score → admission gate → persist (chunks + scores + entity index) +//! The hot path now does: +//! `canonicalise -> chunk -> fast score -> persist chunks/score rows -> enqueue extract jobs` //! -//! Phase 2 inserts scoring between chunker and persistence. Low-scoring -//! chunks are dropped (their rationale is still persisted to -//! `mem_tree_score` for diagnostics); surviving chunks get their entities -//! indexed so later phases can resolve "which chunks mention Alice?" in -//! O(lookup). +//! The slower work (full extraction, admission, tree buffering, sealing, +//! topic routing, daily digests) runs out of the SQLite-backed jobs queue. -use anyhow::{Context, Result}; +use anyhow::Result; use serde::{Deserialize, Serialize}; use crate::openhuman::config::Config; @@ -19,25 +17,20 @@ use crate::openhuman::memory::tree::canonicalize::{ CanonicalisedSource, }; use crate::openhuman::memory::tree::chunker::{chunk_markdown, ChunkerInput, ChunkerOptions}; -use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked}; +use crate::openhuman::memory::tree::jobs::{self, ExtractChunkPayload, NewJob}; use crate::openhuman::memory::tree::score::{self, ScoreResult, ScoringConfig}; -use crate::openhuman::memory::tree::source_tree::{ - append_leaf, get_or_create_source_tree, LeafRef, -}; use crate::openhuman::memory::tree::store; -use crate::openhuman::memory::tree::topic_tree::route_leaf_to_topic_trees; -use crate::openhuman::memory::tree::types::Chunk; -/// Outcome of one ingest call — extended with per-chunk admission info. +/// Outcome of one ingest call. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct IngestResult { pub source_id: String, - /// Number of chunks that passed the admission gate and were persisted. + /// Number of chunks persisted and queued for async extraction. pub chunks_written: usize, - /// Number of chunks that failed the admission gate and were NOT persisted - /// (their score rationale IS persisted for diagnostics). + /// Number of chunks the cheap fast-score path would drop. Final admission + /// still happens later in the extract job. pub chunks_dropped: usize, - /// IDs of all chunks that were persisted (in source order). + /// IDs of all chunks written and queued. pub chunk_ids: Vec, } @@ -52,7 +45,6 @@ impl IngestResult { } } -/// Ingest a batch of chat messages scoped to one channel/group. pub async fn ingest_chat( config: &Config, source_id: &str, @@ -60,11 +52,6 @@ pub async fn ingest_chat( tags: Vec, batch: ChatBatch, ) -> Result { - log::debug!( - "[memory_tree::ingest] chat source_id={} msg_count={}", - source_id, - batch.messages.len() - ); let canonical = match chat::canonicalise(source_id, owner, &tags, batch).map_err(anyhow::Error::msg)? { Some(c) => c, @@ -73,7 +60,6 @@ pub async fn ingest_chat( persist(config, source_id, canonical).await } -/// Ingest a single email thread. pub async fn ingest_email( config: &Config, source_id: &str, @@ -81,11 +67,6 @@ pub async fn ingest_email( tags: Vec, thread: EmailThread, ) -> Result { - log::debug!( - "[memory_tree::ingest] email source_id={} msg_count={}", - source_id, - thread.messages.len() - ); let canonical = match email::canonicalise(source_id, owner, &tags, thread).map_err(anyhow::Error::msg)? { Some(c) => c, @@ -94,7 +75,6 @@ pub async fn ingest_email( persist(config, source_id, canonical).await } -/// Ingest a single standalone document. pub async fn ingest_document( config: &Config, source_id: &str, @@ -102,13 +82,6 @@ pub async fn ingest_document( tags: Vec, doc: DocumentInput, ) -> Result { - let title_len = doc.title.chars().count(); - log::debug!( - "[memory_tree::ingest] document source_id={} has_title={} title_len={}", - source_id, - !doc.title.trim().is_empty(), - title_len - ); let canonical = match document::canonicalise(source_id, owner, &tags, doc).map_err(anyhow::Error::msg)? { Some(c) => c, @@ -122,7 +95,6 @@ async fn persist( source_id: &str, canonical: CanonicalisedSource, ) -> Result { - // 1. Chunk let input = ChunkerInput { source_kind: canonical.metadata.source_kind, source_id: source_id.to_string(), @@ -134,15 +106,8 @@ async fn persist( return Ok(IngestResult::empty(source_id)); } - // 2. Score (async; extractor chain driven by config.memory_tree — - // regex-only unless llm_extractor_endpoint + model are set, in - // which case LlmEntityExtractor runs as second-pass on borderline - // chunks with soft-fallback to regex on LLM failure). let scoring_cfg = ScoringConfig::from_config(config); - let scores = score::score_chunks(&chunks, &scoring_cfg).await?; - - // Fail fast on scorer length mismatch — silently truncating via zip would - // drop chunks (or their score rationale) without trace. + let scores = score::score_chunks_fast(&chunks, &scoring_cfg).await?; if scores.len() != chunks.len() { anyhow::bail!( "[memory_tree::ingest] scorer length mismatch: chunks={} scores={}", @@ -151,78 +116,79 @@ async fn persist( ); } - // 3. Partition kept vs dropped - let mut kept_chunks: Vec = Vec::new(); - let mut all_results: Vec<(ScoreResult, i64)> = Vec::new(); - for (chunk, result) in chunks.iter().zip(scores.into_iter()) { - let ts_ms = chunk.metadata.timestamp.timestamp_millis(); - if result.kept { - kept_chunks.push(chunk.clone()); - } - all_results.push((result, ts_ms)); - } - + let all_results: Vec<(ScoreResult, i64)> = chunks + .iter() + .zip(scores.into_iter()) + .map(|(chunk, result)| (result, chunk.metadata.timestamp.timestamp_millis())) + .collect(); let dropped = all_results.iter().filter(|(r, _)| !r.kept).count(); - log::debug!( - "[memory_tree::ingest] scoring source_id={} kept={} dropped={}", - source_id, - kept_chunks.len(), - dropped - ); - // 3.5. Phase 4 (#710) — embed every kept chunk BEFORE writing. A failed - // embed aborts the whole batch so a retry stays idempotent on - // `chunk_id` (upsert_chunks_tx ON CONFLICT REPLACE). Serial per - // chunk for now; parallelism is an explicit follow-up. We pack - // each vector into its SQLite BLOB representation up front so the - // write tx below doesn't have to juggle floats. - let embedder = build_embedder_from_config(config).context("build embedder during ingest")?; - log::debug!( - "[memory_tree::ingest] embedding source_id={} provider={} kept={}", - source_id, - embedder.name(), - kept_chunks.len() - ); - let mut chunk_embedding_blobs: Vec<(String, Vec)> = Vec::with_capacity(kept_chunks.len()); - for chunk in &kept_chunks { - let vector = embedder - .embed(&chunk.content) - .await - .with_context(|| format!("embed chunk_id={} during ingest", chunk.id))?; - let packed = pack_checked(&vector) - .with_context(|| format!("pack embedding for chunk_id={} during ingest", chunk.id))?; - chunk_embedding_blobs.push((chunk.id.clone(), packed)); - } - - // 4. Persist (blocking SQLite — isolate on a dedicated thread). - // Chunks + scores + embeddings all commit in one tx so retries are - // idempotent. The chunk upsert preserves any prior embedding on - // conflict (see store::upsert_chunks_tx), so we follow with a - // deliberate UPDATE that writes the fresh blob. let config_owned = config.clone(); - let kept_for_store = kept_chunks.clone(); + let chunks_for_store = chunks.clone(); let results_for_store = all_results.clone(); - let embeddings_for_store = chunk_embedding_blobs.clone(); let written = tokio::task::spawn_blocking(move || -> Result { + use std::collections::{HashMap, HashSet}; store::with_connection(&config_owned, |conn| { let tx = conn.unchecked_transaction()?; - let n = store::upsert_chunks_tx(&tx, &kept_for_store)?; - for (chunk_id, blob) in &embeddings_for_store { - let changed = tx.execute( - "UPDATE mem_tree_chunks SET embedding = ?1 WHERE id = ?2", - rusqlite::params![blob, chunk_id], - )?; - if changed == 0 { - log::warn!( - "[memory_tree::ingest] embedding update affected 0 rows chunk_id={chunk_id} — \ - upsert missed the row?" - ); + + // Read each chunk's CURRENT lifecycle BEFORE the upsert. This + // is the "did this chunk exist before this batch" snapshot, + // because `upsert_chunks_tx` will either preserve the existing + // row's lifecycle (UPDATE doesn't touch the column) or insert + // a new row that picks up the column DEFAULT — so reading + // post-upsert can't distinguish "brand new" from + // "already-admitted-from-prior-ingest". + let mut prior: HashMap> = HashMap::new(); + for chunk in &chunks_for_store { + let status = store::get_chunk_lifecycle_status_tx(&tx, &chunk.id)?; + prior.insert(chunk.id.clone(), status); + } + + let n = store::upsert_chunks_tx(&tx, &chunks_for_store)?; + + // Re-ingest of identical content (same chunk_id) must NOT + // downgrade chunks that have already progressed through the + // async pipeline. Without this guard, a re-ingest would reset + // every chunk to 'pending_extraction' and enqueue a fresh + // `extract_chunk` job — sending already-buffered/sealed + // chunks back through extract → admit → append, ultimately + // duplicating them into a second summary in the same tree. + // + // Schedule a chunk for processing when its PRE-upsert state + // was either absent (genuinely new) or already + // `pending_extraction` (a prior ingest crashed before extract + // ran). Anything else — `admitted`, `buffered`, `sealed`, + // `dropped` — is past the point of accepting new work, so + // leave the lifecycle alone and skip the extract enqueue. + let mut to_schedule: HashSet = HashSet::new(); + for chunk in &chunks_for_store { + let pre = prior.get(&chunk.id).cloned().flatten(); + let needs_processing = matches!( + pre.as_deref(), + None | Some(store::CHUNK_STATUS_PENDING_EXTRACTION), + ); + if needs_processing { + store::set_chunk_lifecycle_status_tx( + &tx, + &chunk.id, + store::CHUNK_STATUS_PENDING_EXTRACTION, + )?; + to_schedule.insert(chunk.id.clone()); } } + for (result, ts_ms) in &results_for_store { - // Persist rationale for EVERY chunk (kept or dropped). - // Index entities only for kept chunks (handled inside persist_score_tx). + if !to_schedule.contains(&result.chunk_id) { + // Chunk has already progressed past pending_extraction + // on a prior ingest — skip score re-persist and don't + // enqueue a duplicate extract job. + continue; + } score::persist_score_tx(&tx, result, *ts_ms, None)?; + let extract = NewJob::extract_chunk(&ExtractChunkPayload { + chunk_id: result.chunk_id.clone(), + })?; + let _ = jobs::enqueue_tx(&tx, &extract)?; } tx.commit()?; Ok(n) @@ -231,98 +197,26 @@ async fn persist( .await .map_err(|e| anyhow::anyhow!("persist join error: {e}"))??; - // 5. Source-tree append (Phase 3a #709). Each kept leaf pushes into - // the tree's L0 buffer and cascades upward when token_sum crosses - // the budget. Entities/topics from the scorer are threaded in so - // sealed summaries inherit the child signal set. Failures here - // log at warn level but don't fail the ingest — leaves are already - // persisted, and a later flush/retry can still rebuild the tree. - if let Err(e) = append_leaves_to_tree(config, source_id, &kept_chunks, &all_results).await { - log::warn!( - "[memory_tree::ingest] source_tree append failed source_id={} err={:#}", - source_id, - e - ); - } + jobs::wake_workers(); Ok(IngestResult { source_id: source_id.to_string(), chunks_written: written, chunks_dropped: dropped, - chunk_ids: kept_chunks.iter().map(|c| c.id.clone()).collect(), + chunk_ids: chunks.iter().map(|c| c.id.clone()).collect(), }) } -/// Push every kept chunk into its source tree. Scoped to Phase 3a — all -/// chunks from one ingest batch share the same `source_id`, so they share -/// one tree lookup. The summariser comes from `build_summariser(config)`: -/// Ollama-backed [`LlmSummariser`] when `memory_tree.llm_summariser_*` is -/// set, else deterministic [`InertSummariser`]. Both share the same trait, -/// so the seal cascade doesn't notice the difference. -async fn append_leaves_to_tree( - config: &Config, - source_id: &str, - kept_chunks: &[Chunk], - all_results: &[(ScoreResult, i64)], -) -> Result<()> { - if kept_chunks.is_empty() { - return Ok(()); - } - let tree = get_or_create_source_tree(config, source_id)?; - let summariser = crate::openhuman::memory::tree::source_tree::build_summariser(config); - - // Build a chunk_id → (score, entities, topics) map for quick lookup. - use std::collections::HashMap; - let mut score_by_id: HashMap = HashMap::new(); - for (r, _) in all_results { - score_by_id.insert(r.chunk_id.clone(), r); - } - - for chunk in kept_chunks { - let (score_value, entities, topics) = match score_by_id.get(&chunk.id) { - Some(r) => ( - r.total, - r.canonical_entities - .iter() - .map(|e| e.canonical_id.clone()) - .collect::>(), - chunk.metadata.tags.clone(), - ), - None => (0.0, Vec::new(), chunk.metadata.tags.clone()), - }; - let leaf = LeafRef { - chunk_id: chunk.id.clone(), - token_count: chunk.token_count, - timestamp: chunk.metadata.timestamp, - content: chunk.content.clone(), - entities: entities.clone(), - topics, - score: score_value, - }; - append_leaf(config, &tree, &leaf, summariser.as_ref()).await?; - - // Phase 3c (#709): route the leaf to every matching topic tree - // and tick the curator for each entity. Non-fatal on error — - // the source-tree append has already succeeded above. - if let Err(e) = - route_leaf_to_topic_trees(config, &leaf, &entities, summariser.as_ref()).await - { - log::warn!( - "[memory_tree::ingest] topic_tree routing failed chunk_id={} err={:#}", - chunk.id, - e - ); - } - } - Ok(()) -} - #[cfg(test)] mod tests { use super::*; use crate::openhuman::memory::tree::canonicalize::chat::ChatMessage; + use crate::openhuman::memory::tree::jobs::drain_until_idle; use crate::openhuman::memory::tree::score::store::{count_scores, lookup_entity}; - use crate::openhuman::memory::tree::store::{count_chunks, list_chunks, ListChunksQuery}; + use crate::openhuman::memory::tree::store::{ + count_chunks, count_chunks_by_lifecycle_status, get_chunk_embedding, list_chunks, + ListChunksQuery, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, + }; use crate::openhuman::memory::tree::types::SourceKind; use chrono::{TimeZone, Utc}; use tempfile::TempDir; @@ -331,16 +225,12 @@ mod tests { let tmp = TempDir::new().unwrap(); let mut cfg = Config::default(); cfg.workspace_dir = tmp.path().to_path_buf(); - // Phase 4 (#710): disable Ollama-backed embeddings in tests. - // Ingest/seal call `build_embedder_from_config`; falling back to - // inert (zero vectors) keeps tests deterministic and network-free. cfg.memory_tree.embedding_endpoint = None; cfg.memory_tree.embedding_model = None; cfg.memory_tree.embedding_strict = false; (tmp, cfg) } - /// Build a substantive batch that reliably passes the admission gate. fn substantive_batch() -> ChatBatch { ChatBatch { platform: "slack".into(), @@ -349,17 +239,14 @@ mod tests { ChatMessage { author: "alice".into(), timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: "We are planning to ship the Phoenix migration on Friday \ - after reviewing the runbook and staging results. Please \ - confirm availability by replying here. alice@example.com" + text: "We are planning to ship the Phoenix migration on Friday after reviewing the runbook and staging results. alice@example.com" .into(), source_ref: Some("slack://m1".into()), }, ChatMessage { author: "bob".into(), timestamp: Utc.timestamp_millis_opt(1_700_000_010_000).unwrap(), - text: "Confirmed — I'll handle the coordination and cut a release \ - candidate tonight. #launch-q2 will be tracked in Notion." + text: "Confirmed, I will handle the coordination and launch tracking tonight." .into(), source_ref: None, }, @@ -368,25 +255,38 @@ mod tests { } #[tokio::test] - async fn ingest_chat_writes_substantive_chunks() { + async fn ingest_chat_writes_and_queue_drains_to_admitted_chunk() { let (_tmp, cfg) = test_config(); let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()) .await .unwrap(); assert_eq!(out.chunks_written, 1); - assert_eq!(out.chunks_dropped, 0); assert_eq!(count_chunks(&cfg).unwrap(), 1); - // Score row persisted for the kept chunk + + drain_until_idle(&cfg).await.unwrap(); + + // Final lifecycle is `buffered`: extract → admitted → append_buffer → buffered. + // The single small chunk doesn't cross TOKEN_BUDGET so no seal fires. + assert_eq!( + count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_BUFFERED).unwrap(), + 1 + ); assert_eq!(count_scores(&cfg).unwrap(), 1); - // Entity index populated from regex extraction (alice@example.com + hashtag) - let alice_hits = lookup_entity(&cfg, "email:alice@example.com", None).unwrap(); - assert_eq!(alice_hits.len(), 1); + assert_eq!( + lookup_entity(&cfg, "email:alice@example.com", None) + .unwrap() + .len(), + 1 + ); let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); assert_eq!(rows[0].metadata.source_kind, SourceKind::Chat); + assert!(get_chunk_embedding(&cfg, &out.chunk_ids[0]) + .unwrap() + .is_some()); } #[tokio::test] - async fn low_signal_chunks_are_dropped_but_score_persists() { + async fn low_signal_chunks_end_up_dropped_after_queue_processing() { let (_tmp, cfg) = test_config(); let batch = ChatBatch { platform: "slack".into(), @@ -394,18 +294,22 @@ mod tests { messages: vec![ChatMessage { author: "alice".into(), timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), - text: "+1".into(), // extremely low-signal + text: "+1".into(), source_ref: None, }], }; let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch) .await .unwrap(); - assert_eq!(out.chunks_written, 0); - assert_eq!(out.chunks_dropped, 1); - // Chunk NOT in chunks table - assert_eq!(count_chunks(&cfg).unwrap(), 0); - // Score row IS persisted for diagnostics + assert_eq!(out.chunks_written, 1); + assert_eq!(count_chunks(&cfg).unwrap(), 1); + + drain_until_idle(&cfg).await.unwrap(); + + assert_eq!( + count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_DROPPED).unwrap(), + 1 + ); assert_eq!(count_scores(&cfg).unwrap(), 1); } @@ -421,7 +325,6 @@ mod tests { .await .unwrap(); assert_eq!(out.chunks_written, 0); - assert_eq!(out.chunks_dropped, 0); assert_eq!(count_chunks(&cfg).unwrap(), 0); assert_eq!(count_scores(&cfg).unwrap(), 0); } @@ -432,9 +335,7 @@ mod tests { let doc = DocumentInput { provider: "notion".into(), title: "Launch plan".into(), - body: "We are planning to ship Phoenix on Friday after review. \ - Coordination is via email and the launch thread tracks \ - the relevant decisions. alice@example.com owns this." + body: "We are planning to ship Phoenix on Friday after review. alice@example.com owns this." .into(), modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), source_ref: Some("notion://page/abc".into()), @@ -445,74 +346,8 @@ mod tests { ingest_document(&cfg, "notion:abc", "alice", vec![], doc) .await .unwrap(); + drain_until_idle(&cfg).await.unwrap(); assert_eq!(count_chunks(&cfg).unwrap(), 1); assert_eq!(count_scores(&cfg).unwrap(), 1); } - - #[tokio::test] - async fn chunks_preserve_source_ref_when_kept() { - let (_tmp, cfg) = test_config(); - let doc = DocumentInput { - provider: "notion".into(), - title: "t".into(), - body: "Phoenix launch plan with enough substance to pass the admission \ - gate: we are reviewing the migration runbook alice@example.com \ - on Friday evening." - .into(), - modified_at: Utc::now(), - source_ref: Some("notion://x".into()), - }; - ingest_document(&cfg, "notion:x", "alice", vec![], doc) - .await - .unwrap(); - let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); - assert_eq!(rows.len(), 1); - assert_eq!( - rows[0].metadata.source_ref.as_ref().unwrap().value, - "notion://x" - ); - } - - // ── Phase 4 (#710) ────────────────────────────────────────────── - - /// Ingesting with an inert embedder must still write embeddings for - /// every kept chunk. Use a trimmed test_config so we go through the - /// full embed → persist pipeline. - #[tokio::test] - async fn ingest_writes_chunk_embeddings() { - use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM; - use crate::openhuman::memory::tree::store::get_chunk_embedding; - - let (_tmp, cfg) = test_config(); - let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()) - .await - .unwrap(); - assert!(out.chunks_written >= 1); - for id in &out.chunk_ids { - let v = get_chunk_embedding(&cfg, id) - .unwrap() - .unwrap_or_else(|| panic!("missing embedding for {id}")); - assert_eq!(v.len(), EMBEDDING_DIM); - } - } - - /// When the embedder errors (here: Ollama required but unavailable), - /// ingest must fail and persist nothing. Point embedding at a dead - /// port so the HTTP call refuses, and flip `strict=true` so the - /// factory returns Ollama (not the inert fallback). - #[tokio::test] - async fn ingest_fails_when_embedder_fails_and_persists_nothing() { - let (_tmp, mut cfg) = test_config(); - cfg.memory_tree.embedding_endpoint = Some("http://127.0.0.1:1".into()); - cfg.memory_tree.embedding_model = Some("nomic-embed-text".into()); - cfg.memory_tree.embedding_timeout_ms = Some(500); - cfg.memory_tree.embedding_strict = true; - - let result = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()).await; - assert!(result.is_err(), "expected ingest to bail when embed fails"); - - // No chunks, no scores persisted — retry stays clean. - assert_eq!(count_chunks(&cfg).unwrap(), 0); - assert_eq!(count_scores(&cfg).unwrap(), 0); - } } diff --git a/src/openhuman/memory/tree/jobs/handlers/mod.rs b/src/openhuman/memory/tree/jobs/handlers/mod.rs new file mode 100644 index 00000000000..49d7a5d2b4d --- /dev/null +++ b/src/openhuman/memory/tree/jobs/handlers/mod.rs @@ -0,0 +1,719 @@ +use anyhow::{Context, Result}; +use chrono::TimeZone; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::global_tree::digest::{self, DigestOutcome}; +use crate::openhuman::memory::tree::jobs::store; +use crate::openhuman::memory::tree::jobs::types::{ + AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload, + Job, JobKind, NewJob, NodeRef, SealPayload, TopicRoutePayload, +}; +use crate::openhuman::memory::tree::score; +use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked}; +use crate::openhuman::memory::tree::score::extract::build_summary_extractor; +use crate::openhuman::memory::tree::score::store as score_store; +use crate::openhuman::memory::tree::source_tree::{ + append_leaf_deferred, build_summariser, get_or_create_source_tree, LabelStrategy, LeafRef, +}; +use crate::openhuman::memory::tree::store as chunk_store; +use crate::openhuman::memory::tree::topic_tree::curator; + +pub async fn handle_job(config: &Config, job: &Job) -> Result<()> { + match job.kind { + JobKind::ExtractChunk => handle_extract(config, job).await, + JobKind::AppendBuffer => handle_append_buffer(config, job).await, + JobKind::Seal => handle_seal(config, job).await, + JobKind::TopicRoute => handle_topic_route(config, job).await, + JobKind::DigestDaily => handle_digest_daily(config, job).await, + JobKind::FlushStale => handle_flush_stale(config, job).await, + } +} + +async fn handle_extract(config: &Config, job: &Job) -> Result<()> { + let payload: ExtractChunkPayload = + serde_json::from_str(&job.payload_json).context("parse ExtractChunk payload")?; + let Some(chunk) = chunk_store::get_chunk(config, &payload.chunk_id)? else { + log::warn!( + "[memory_tree::jobs] extract chunk missing chunk_id={}", + payload.chunk_id + ); + return Ok(()); + }; + + let scoring_cfg = score::ScoringConfig::from_config(config); + let result = score::score_chunk(&chunk, &scoring_cfg).await?; + let packed_embedding = if result.kept { + let embedder = + build_embedder_from_config(config).context("build embedder in extract handler")?; + let vector = embedder + .embed(&chunk.content) + .await + .with_context(|| format!("embed chunk_id={} in extract handler", chunk.id))?; + Some( + pack_checked(&vector) + .with_context(|| format!("pack embedding for chunk_id={}", chunk.id))?, + ) + } else { + None + }; + + chunk_store::with_connection(config, |conn| { + let tx = conn.unchecked_transaction()?; + score::persist_score_tx( + &tx, + &result, + chunk.metadata.timestamp.timestamp_millis(), + None, + )?; + + if result.kept { + tx.execute( + "UPDATE mem_tree_chunks + SET embedding = ?1, + lifecycle_status = ?2 + WHERE id = ?3", + rusqlite::params![ + packed_embedding, + chunk_store::CHUNK_STATUS_ADMITTED, + chunk.id, + ], + )?; + } else { + tx.execute( + "UPDATE mem_tree_chunks + SET lifecycle_status = ?1 + WHERE id = ?2", + rusqlite::params![chunk_store::CHUNK_STATUS_DROPPED, chunk.id], + )?; + } + + tx.commit()?; + Ok(()) + })?; + + if !result.kept { + return Ok(()); + } + + let source_job = NewJob::append_buffer(&AppendBufferPayload { + node: NodeRef::Leaf { + chunk_id: chunk.id.clone(), + }, + target: AppendTarget::Source { + source_id: chunk.metadata.source_id.clone(), + }, + })?; + if store::enqueue(config, &source_job)?.is_some() { + super::worker::wake_workers(); + } + + let route_job = NewJob::topic_route(&TopicRoutePayload { + node: NodeRef::Leaf { + chunk_id: chunk.id.clone(), + }, + })?; + if store::enqueue(config, &route_job)?.is_some() { + super::worker::wake_workers(); + } + + Ok(()) +} + +async fn handle_append_buffer(config: &Config, job: &Job) -> Result<()> { + use crate::openhuman::memory::tree::source_tree::bucket_seal::should_seal; + use crate::openhuman::memory::tree::source_tree::store as src_store; + + let payload: AppendBufferPayload = + serde_json::from_str(&job.payload_json).context("parse AppendBuffer payload")?; + + // Hydrate the leaf-shaped record from either a chunk row or a summary + // row. The downstream buffer-push doesn't care which kind produced + // the LeafRef. + let (leaf, chunk_id_for_lifecycle): (LeafRef, Option) = match &payload.node { + NodeRef::Leaf { chunk_id } => { + let Some(chunk) = chunk_store::get_chunk(config, chunk_id)? else { + log::warn!("[memory_tree::jobs] append_buffer chunk missing chunk_id={chunk_id}"); + return Ok(()); + }; + let score_row = score_store::get_score(config, &chunk.id)? + .ok_or_else(|| anyhow::anyhow!("missing score row for chunk {}", chunk.id))?; + let entity_ids = score_store::list_entity_ids_for_node(config, &chunk.id)?; + let leaf = LeafRef { + chunk_id: chunk.id.clone(), + token_count: chunk.token_count, + timestamp: chunk.metadata.timestamp, + content: chunk.content.clone(), + entities: entity_ids, + topics: chunk.metadata.tags.clone(), + score: score_row.total, + }; + (leaf, Some(chunk.id)) + } + NodeRef::Summary { summary_id } => { + let Some(summary) = src_store::get_summary(config, summary_id)? else { + log::warn!( + "[memory_tree::jobs] append_buffer summary missing summary_id={summary_id}" + ); + return Ok(()); + }; + // Build a LeafRef from the summary's already-populated fields. + // `chunk_id` carries the source-node id (any string); buffer + // accounting uses it as the item id only. + let leaf = LeafRef { + chunk_id: summary.id.clone(), + token_count: summary.token_count, + timestamp: summary.time_range_start, + content: summary.content.clone(), + entities: summary.entities.clone(), + topics: summary.topics.clone(), + score: summary.score, + }; + (leaf, None) // summaries have no chunk lifecycle to update + } + }; + + // Resolve target tree (no tx open yet — this can create a row). + let tree = match &payload.target { + AppendTarget::Source { source_id } => Some(get_or_create_source_tree(config, source_id)?), + AppendTarget::Topic { tree_id } => src_store::get_tree(config, tree_id)?, + }; + let Some(tree) = tree else { + // Target topic tree doesn't exist (e.g. archived between + // topic_route and this append). Drop on the floor — the + // topic_route was advisory and the source-tree path already + // ran for this leaf. + return Ok(()); + }; + + let is_source_target = matches!(payload.target, AppendTarget::Source { .. }); + let leaf_for_tx = leaf.clone(); + let tree_for_tx = tree.clone(); + let lifecycle_chunk_id = chunk_id_for_lifecycle.clone(); + + // ATOMIC: buffer push + seal enqueue (if gate met) + lifecycle update + // happen in a single SQLite transaction. Eliminates the crash window + // where the buffer commits but the seal job is lost — which can + // duplicate the leaf into two summaries on retry-after-seal-cleared. + let did_enqueue_seal = chunk_store::with_connection(config, move |conn| { + let tx = conn.unchecked_transaction()?; + + // 1. Push leaf into L0 buffer (idempotent on (tree, level, item_id)). + let mut buf = src_store::get_buffer_conn(&tx, &tree_for_tx.id, 0)?; + if !buf.item_ids.iter().any(|x| x == &leaf_for_tx.chunk_id) { + buf.item_ids.push(leaf_for_tx.chunk_id.clone()); + buf.token_sum = buf.token_sum.saturating_add(leaf_for_tx.token_count as i64); + buf.oldest_at = match buf.oldest_at { + Some(existing) => Some(existing.min(leaf_for_tx.timestamp)), + None => Some(leaf_for_tx.timestamp), + }; + src_store::upsert_buffer_tx(&tx, &buf)?; + } + + // 2. If the gate is met, enqueue a seal job atomically. + let did_enqueue = if should_seal(&buf) { + let seal = SealPayload { + tree_id: tree_for_tx.id.clone(), + level: 0, + force_now_ms: None, + }; + store::enqueue_tx(&tx, &NewJob::seal(&seal)?)?.is_some() + } else { + false + }; + + // 3. Lifecycle transition (Source target with a leaf chunk). + // Last step in the tx — its presence is the "this handler + // finished" marker. Same tx as the push + seal-enqueue, so a + // crash anywhere rolls everything back together. + if is_source_target { + if let Some(chunk_id) = lifecycle_chunk_id.as_deref() { + chunk_store::set_chunk_lifecycle_status_tx( + &tx, + chunk_id, + chunk_store::CHUNK_STATUS_BUFFERED, + )?; + } + } + + tx.commit()?; + Ok(did_enqueue) + })?; + + if did_enqueue_seal { + super::worker::wake_workers(); + } + Ok(()) +} + +async fn handle_seal(config: &Config, job: &Job) -> Result<()> { + use crate::openhuman::memory::tree::source_tree::bucket_seal::{seal_one_level, should_seal}; + use crate::openhuman::memory::tree::source_tree::store as src_store; + use crate::openhuman::memory::tree::source_tree::types::TreeKind; + + let payload: SealPayload = + serde_json::from_str(&job.payload_json).context("parse Seal payload")?; + let Some(tree) = src_store::get_tree(config, &payload.tree_id)? else { + log::warn!( + "[memory_tree::jobs] seal tree missing tree_id={}", + payload.tree_id + ); + return Ok(()); + }; + + // Seal exactly one level. Parents only get sealed via a follow-up job + // so each level is its own crash-recovery checkpoint and each LLM + // summariser call competes for a fresh slot from the global semaphore. + let buf = src_store::get_buffer(config, &tree.id, payload.level)?; + let forced = payload.force_now_ms.is_some(); + if buf.is_empty() { + log::debug!( + "[memory_tree::jobs] seal skipped — empty buffer tree_id={} level={}", + tree.id, + payload.level + ); + return Ok(()); + } + if !forced && !should_seal(&buf) { + // Another job sealed this level out from under us (or the buffer + // hasn't crossed the gate yet); idempotent no-op. + log::debug!( + "[memory_tree::jobs] seal gate not met tree_id={} level={} token_sum={}", + tree.id, + payload.level, + buf.token_sum + ); + return Ok(()); + } + + // Pick the labeling strategy for this tree kind. Source trees mint + // emergent themes via the seal-time extractor; topic trees stay empty + // by design (scope already pins the canonical id). Global trees never + // reach here — `digest_daily` handles them — but Empty is a safe + // defensive default. + let strategy = match tree.kind { + TreeKind::Source => LabelStrategy::ExtractFromContent(build_summary_extractor(config)), + TreeKind::Topic => LabelStrategy::Empty, + TreeKind::Global => LabelStrategy::Empty, + }; + + let summariser = build_summariser(config); + // `seal_one_level` with `enqueue_follow_ups: true` atomically inserts + // the parent-cascade seal (if the parent buffer now meets its gate) + // and the summary-side `topic_route` (for source trees) inside the + // same SQLite transaction that commits the seal. This eliminates the + // crash window where the seal succeeds but the follow-up enqueues + // are silently lost. + let _summary_id = + seal_one_level(config, &tree, &buf, summariser.as_ref(), &strategy, true).await?; + super::worker::wake_workers(); + Ok(()) +} + +async fn handle_topic_route(config: &Config, job: &Job) -> Result<()> { + let payload: TopicRoutePayload = + serde_json::from_str(&job.payload_json).context("parse TopicRoute payload")?; + + // Resolve the source node id and verify it exists. `mem_tree_entity_index` + // already indexes both chunks and summaries via `node_kind`, so the + // canonical-id loop below is identical for either case. + let node_id: String = match &payload.node { + NodeRef::Leaf { chunk_id } => { + if chunk_store::get_chunk(config, chunk_id)?.is_none() { + log::warn!("[memory_tree::jobs] topic_route chunk missing chunk_id={chunk_id}"); + return Ok(()); + } + chunk_id.clone() + } + NodeRef::Summary { summary_id } => { + if crate::openhuman::memory::tree::source_tree::store::get_summary(config, summary_id)? + .is_none() + { + log::warn!( + "[memory_tree::jobs] topic_route summary missing summary_id={summary_id}" + ); + return Ok(()); + } + summary_id.clone() + } + }; + + let entity_ids = score_store::list_entity_ids_for_node(config, &node_id)?; + if entity_ids.is_empty() { + log::debug!("[memory_tree::jobs] topic_route no entities for node_id={node_id} — skipping"); + return Ok(()); + } + + let summariser = build_summariser(config); + for entity_id in entity_ids { + let _ = curator::maybe_spawn_topic_tree(config, &entity_id, summariser.as_ref()).await?; + if let Some(tree) = crate::openhuman::memory::tree::source_tree::store::get_tree_by_scope( + config, + crate::openhuman::memory::tree::source_tree::types::TreeKind::Topic, + &entity_id, + )? { + let job = NewJob::append_buffer(&AppendBufferPayload { + node: payload.node.clone(), + target: AppendTarget::Topic { + tree_id: tree.id.clone(), + }, + })?; + if store::enqueue(config, &job)?.is_some() { + super::worker::wake_workers(); + } + } + } + Ok(()) +} + +async fn handle_digest_daily(config: &Config, job: &Job) -> Result<()> { + let payload: DigestDailyPayload = + serde_json::from_str(&job.payload_json).context("parse DigestDaily payload")?; + let day = chrono::NaiveDate::parse_from_str(&payload.date_iso, "%Y-%m-%d") + .with_context(|| format!("invalid digest date {}", payload.date_iso))?; + let summariser = build_summariser(config); + match digest::end_of_day_digest(config, day, summariser.as_ref()).await? { + DigestOutcome::Emitted { daily_id, .. } => { + log::info!("[memory_tree::jobs] emitted digest daily_id={daily_id}"); + } + DigestOutcome::EmptyDay => {} + DigestOutcome::Skipped { existing_id } => { + log::debug!("[memory_tree::jobs] digest skipped existing_id={existing_id}"); + } + } + Ok(()) +} + +async fn handle_flush_stale(config: &Config, job: &Job) -> Result<()> { + let payload: FlushStalePayload = + serde_json::from_str(&job.payload_json).context("parse FlushStale payload")?; + let age_secs = payload + .max_age_secs + .unwrap_or(crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS); + let cutoff = chrono::Utc::now() - chrono::Duration::seconds(age_secs); + let buffers = + crate::openhuman::memory::tree::source_tree::store::list_stale_buffers(config, cutoff)?; + for buf in buffers { + let seal = SealPayload { + tree_id: buf.tree_id.clone(), + level: buf.level, + force_now_ms: Some(chrono::Utc::now().timestamp_millis()), + }; + if store::enqueue(config, &NewJob::seal(&seal)?)?.is_some() { + super::worker::wake_workers(); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::jobs::store::{count_by_status, count_total}; + use crate::openhuman::memory::tree::jobs::types::JobStatus; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf_deferred, LeafRef}; + use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; + use crate::openhuman::memory::tree::source_tree::store as src_store; + use crate::openhuman::memory::tree::source_tree::types::TreeKind; + use crate::openhuman::memory::tree::store::with_connection; + use chrono::TimeZone; + use rusqlite::params; + use tempfile::TempDir; + + fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + (tmp, cfg) + } + + /// Build a minimal `Job` row for direct handler invocation. Mirrors + /// what `claim_next` would produce for a freshly-claimed row. + fn mk_running_job(kind: JobKind, payload_json: String) -> Job { + let now_ms = chrono::Utc::now().timestamp_millis(); + Job { + id: "test-job-id".into(), + kind, + payload_json, + dedupe_key: None, + status: JobStatus::Running, + attempts: 1, + max_attempts: 5, + available_at_ms: now_ms, + locked_until_ms: Some(now_ms + 60_000), + last_error: None, + created_at_ms: now_ms, + started_at_ms: Some(now_ms), + completed_at_ms: None, + } + } + + /// Count rows in `mem_tree_jobs` matching a specific kind. + fn count_jobs_of_kind(cfg: &Config, kind: &str) -> u64 { + with_connection(cfg, |conn| { + let n: i64 = conn.query_row( + "SELECT COUNT(*) FROM mem_tree_jobs WHERE kind = ?1", + params![kind], + |r| r.get(0), + )?; + Ok(n.max(0) as u64) + }) + .unwrap() + } + + /// Seed a source tree and push enough labeled leaves into its L0 buffer + /// to cross `TOKEN_BUDGET`, returning the tree. The caller can then + /// fire `handle_seal` and inspect the result. + async fn seed_source_tree_ready_to_seal( + cfg: &Config, + ) -> crate::openhuman::memory::tree::source_tree::types::Tree { + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{ + chunk_id, Chunk, Metadata, SourceKind, SourceRef, + }; + let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap(); + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let chunk = Chunk { + id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "handler-seed"), + content: "alice@example.com leading the rollout".into(), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: "slack:#eng".into(), + owner: "alice".into(), + timestamp: ts, + time_range: (ts, ts), + tags: vec![], + source_ref: Some(SourceRef::new("slack://x")), + }, + // Bust budget so the L0 buffer is "ready" for seal. + token_count: 10_000, + seq_in_source: 0, + created_at: ts, + }; + upsert_chunks(cfg, &[chunk.clone()]).unwrap(); + let leaf = LeafRef { + chunk_id: chunk.id, + token_count: 10_000, + timestamp: ts, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 0.5, + }; + // append_leaf_deferred only buffers; doesn't seal. handle_seal will. + let _ = append_leaf_deferred(cfg, &tree, &leaf).unwrap(); + tree + } + + #[tokio::test] + async fn source_tree_seal_handler_enqueues_summary_topic_route() { + let (_tmp, cfg) = test_config(); + let tree = seed_source_tree_ready_to_seal(&cfg).await; + + let payload = SealPayload { + tree_id: tree.id.clone(), + level: 0, + force_now_ms: None, + }; + let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap()); + + // Pre-condition: queue has no topic_route jobs. + assert_eq!(count_jobs_of_kind(&cfg, "topic_route"), 0); + + super::handle_seal(&cfg, &job).await.unwrap(); + + // Post-condition: source-tree seal must enqueue exactly one + // topic_route job carrying NodeRef::Summary { summary_id: }. + assert_eq!( + count_jobs_of_kind(&cfg, "topic_route"), + 1, + "source-tree seal must enqueue summary-side topic_route" + ); + assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1); + + // Inspect the enqueued payload to confirm it's a Summary variant. + let payload_json: String = with_connection(&cfg, |conn| { + let s: String = conn + .query_row( + "SELECT payload_json FROM mem_tree_jobs WHERE kind = 'topic_route'", + [], + |r| r.get(0), + ) + .unwrap(); + Ok(s) + }) + .unwrap(); + let p: TopicRoutePayload = serde_json::from_str(&payload_json).unwrap(); + match p.node { + NodeRef::Summary { summary_id } => { + assert!(summary_id.starts_with("summary:L1:")); + } + other => panic!("expected NodeRef::Summary, got {other:?}"), + } + } + + #[tokio::test] + async fn topic_tree_seal_handler_does_not_enqueue_topic_route() { + let (_tmp, cfg) = test_config(); + // Spawn a topic tree directly via the registry (skipping curator's + // hotness gate — we just need a TreeKind::Topic with leaves). + let topic_tree = + crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree( + &cfg, + "topic:phoenix-migration", + ) + .unwrap(); + // Push a single 10k-token leaf so L0 is gate-ready. + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{ + chunk_id, Chunk, Metadata, SourceKind, SourceRef, + }; + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let chunk = Chunk { + id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "topic-seed"), + content: "topic content".into(), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: "slack:#eng".into(), + owner: "alice".into(), + timestamp: ts, + time_range: (ts, ts), + tags: vec![], + source_ref: Some(SourceRef::new("slack://x")), + }, + token_count: 10_000, + seq_in_source: 0, + created_at: ts, + }; + upsert_chunks(&cfg, &[chunk.clone()]).unwrap(); + let leaf = LeafRef { + chunk_id: chunk.id, + token_count: 10_000, + timestamp: ts, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 0.5, + }; + append_leaf_deferred(&cfg, &topic_tree, &leaf).unwrap(); + + let payload = SealPayload { + tree_id: topic_tree.id.clone(), + level: 0, + force_now_ms: None, + }; + let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap()); + + super::handle_seal(&cfg, &job).await.unwrap(); + + // Topic-tree seals are sinks: must not enqueue any topic_route. + assert_eq!( + count_jobs_of_kind(&cfg, "topic_route"), + 0, + "topic-tree seal must NOT enqueue topic_route (trees are sinks)" + ); + // The seal itself should still have produced a summary node. + assert_eq!(src_store::count_summaries(&cfg, &topic_tree.id).unwrap(), 1); + } + + #[tokio::test] + async fn handle_append_buffer_with_summary_payload_pushes_into_topic_tree() { + let (_tmp, cfg) = test_config(); + + // 1. Create a target topic tree with a clean L0 buffer. + let topic_tree = + crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree( + &cfg, + "email:alice@example.com", + ) + .unwrap(); + let l0_before = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap(); + assert!(l0_before.is_empty()); + + // 2. Manually insert a summary node we can route. The simplest way + // is to create a separate source tree, push two 6k leaves into + // it, and let the seal produce a summary we can address. + let source_tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); + use crate::openhuman::memory::tree::source_tree::bucket_seal::seal_one_level; + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{ + chunk_id, Chunk, Metadata, SourceKind, SourceRef, + }; + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + for seq in 0..2 { + let chunk = Chunk { + id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "summary-seed"), + content: format!("source content {seq}"), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: "slack:#eng".into(), + owner: "alice".into(), + timestamp: ts, + time_range: (ts, ts), + tags: vec![], + source_ref: Some(SourceRef::new("slack://x")), + }, + token_count: 6_000, + seq_in_source: seq, + created_at: ts, + }; + upsert_chunks(&cfg, &[chunk.clone()]).unwrap(); + let leaf = LeafRef { + chunk_id: chunk.id, + token_count: 6_000, + timestamp: ts, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 0.5, + }; + let _ = append_leaf_deferred(&cfg, &source_tree, &leaf).unwrap(); + } + // Force-seal the source tree's L0 to mint the summary. + let buf = src_store::get_buffer(&cfg, &source_tree.id, 0).unwrap(); + let summariser = build_summariser(&cfg); + let summary_id = seal_one_level( + &cfg, + &source_tree, + &buf, + summariser.as_ref(), + &crate::openhuman::memory::tree::source_tree::bucket_seal::LabelStrategy::Empty, + // No follow-up enqueues — the test scopes assertions to the + // append_buffer handler, not seal-side fan-out. + false, + ) + .await + .unwrap(); + + // 3. Build an append_buffer payload routing the summary into the + // topic tree. + let payload = AppendBufferPayload { + node: NodeRef::Summary { + summary_id: summary_id.clone(), + }, + target: AppendTarget::Topic { + tree_id: topic_tree.id.clone(), + }, + }; + let job = mk_running_job( + JobKind::AppendBuffer, + serde_json::to_string(&payload).unwrap(), + ); + + // Clear out any pending append_buffer jobs minted upstream so the + // post-condition assertion below is unambiguous. + let pre = count_total(&cfg).unwrap(); + + super::handle_append_buffer(&cfg, &job).await.unwrap(); + + // 4. Topic tree's L0 buffer should now hold the summary id. + let l0_after = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap(); + assert_eq!(l0_after.item_ids, vec![summary_id]); + assert!(l0_after.token_sum > 0); + + // No new jobs should have been enqueued (buffer didn't cross gate). + assert_eq!(count_total(&cfg).unwrap(), pre); + } +} diff --git a/src/openhuman/memory/tree/jobs/mod.rs b/src/openhuman/memory/tree/jobs/mod.rs new file mode 100644 index 00000000000..bcc5aec9542 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/mod.rs @@ -0,0 +1,45 @@ +//! Async job pipeline for memory-tree work. +//! +//! Replaces the previous synchronous `append_leaf → cascade_seal → LLM +//! summarise` chain on the ingest hot path with a SQLite-backed job queue +//! and a worker pool. The shape is: +//! +//! ```text +//! ingest::persist +//! └── writes chunk row (lifecycle = pending_extraction) +//! enqueues `extract_chunk` +//! +//! worker pool (3 tasks) ──► claims jobs by kind: +//! extract_chunk → LLM extraction → admission decision → enqueue append_buffer +//! append_buffer → push to L0 → enqueue seal if gate met → enqueue topic_route +//! seal → seal one level → enqueue parent seal if cascading +//! topic_route → match topics → enqueue per-topic append_buffer +//! digest_daily → call global_tree::digest::end_of_day_digest +//! flush_stale → enqueue seals for time-stale buffers +//! +//! scheduler (1 task) ──► daily wall-clock tick: +//! enqueues digest_daily(yesterday) + flush_stale(today) +//! ``` +//! +//! All persistence lives in the same `chunks.db` as `mem_tree_chunks` so a +//! producer can insert its side-effect and its follow-up job in one tx. +//! See [`store::enqueue_tx`] for the in-tx producer entry point. + +mod handlers; +pub mod scheduler; +pub mod store; +pub mod testing; +pub mod types; +mod worker; + +pub use scheduler::{backfill_missing_digests, trigger_digest}; +pub use store::{ + claim_next, count_by_status, count_total, enqueue, enqueue_tx, get_job, mark_done, mark_failed, + recover_stale_locks, DEFAULT_LOCK_DURATION_MS, +}; +pub use testing::drain_until_idle; +pub use types::{ + AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload, + Job, JobKind, JobStatus, NewJob, NodeRef, SealPayload, TopicRoutePayload, +}; +pub use worker::{start, wake_workers}; diff --git a/src/openhuman/memory/tree/jobs/scheduler.rs b/src/openhuman/memory/tree/jobs/scheduler.rs new file mode 100644 index 00000000000..9ad4ae691c5 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/scheduler.rs @@ -0,0 +1,211 @@ +use std::time::Duration; + +use anyhow::Result; +use chrono::{Datelike, Duration as ChronoDuration, NaiveDate, TimeZone, Utc}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::jobs::store; +use crate::openhuman::memory::tree::jobs::types::{DigestDailyPayload, FlushStalePayload, NewJob}; + +static STARTED: std::sync::Once = std::sync::Once::new(); + +/// Start the daily wall-clock scheduler. Takes the full `Config` so the +/// digest enqueues match the same workspace + LLM settings the workers +/// see — not `Config::default()`. +pub fn start(config: Config) { + STARTED.call_once(|| { + tokio::spawn(async move { + loop { + if let Err(err) = enqueue_daily_jobs(&config) { + log::warn!("[memory_tree::jobs] scheduler enqueue failed: {err:#}"); + } + tokio::time::sleep(next_sleep_duration()).await; + } + }); + }); +} + +fn enqueue_daily_jobs(config: &Config) -> anyhow::Result<()> { + let now = Utc::now(); + let yesterday = now.date_naive() - ChronoDuration::days(1); + let date_iso = yesterday.format("%Y-%m-%d").to_string(); + + if store::enqueue( + config, + &NewJob::digest_daily(&DigestDailyPayload { + date_iso: date_iso.clone(), + })?, + )? + .is_some() + { + super::worker::wake_workers(); + } + + let today_iso = now.date_naive().format("%Y-%m-%d").to_string(); + if store::enqueue( + config, + &NewJob::flush_stale(&FlushStalePayload::default(), &today_iso)?, + )? + .is_some() + { + super::worker::wake_workers(); + } + + Ok(()) +} + +/// Manually enqueue a `digest_daily` job for `date`. Idempotent — if a +/// digest already ran for that day, the handler's `find_existing_daily` +/// check will return `Skipped` without doing any work; if a job for the +/// same date is already queued or running, the partial unique index on +/// `dedupe_key` suppresses the duplicate. +/// +/// Useful for catch-up after the process was down across midnight, or +/// to force a re-run for testing / debugging. +pub fn trigger_digest(config: &Config, date: NaiveDate) -> Result> { + let payload = DigestDailyPayload { + date_iso: date.format("%Y-%m-%d").to_string(), + }; + let job_id = store::enqueue(config, &NewJob::digest_daily(&payload)?)?; + if job_id.is_some() { + log::info!( + "[memory_tree::jobs] manual digest trigger enqueued date={} id={:?}", + payload.date_iso, + job_id.as_deref() + ); + super::worker::wake_workers(); + } else { + log::debug!( + "[memory_tree::jobs] manual digest trigger dedupe-suppressed date={} \ + (an active job for this date already exists)", + payload.date_iso + ); + } + Ok(job_id) +} + +/// Enqueue `digest_daily` jobs for the last `days_back` calendar days +/// (excluding today). Catch-up helper for cases where the scheduler +/// missed days because the process was down. +/// +/// Returns the number of jobs newly enqueued. Days that already have a +/// completed digest are still re-enqueued — the handler is idempotent +/// and skips them — so this is safe to call repeatedly. +pub fn backfill_missing_digests(config: &Config, days_back: i64) -> Result { + if days_back <= 0 { + return Ok(0); + } + let today = Utc::now().date_naive(); + let mut enqueued = 0usize; + for offset in 1..=days_back { + let date = today - ChronoDuration::days(offset); + if trigger_digest(config, date)?.is_some() { + enqueued += 1; + } + } + log::info!( + "[memory_tree::jobs] backfill_missing_digests window={}d enqueued={}", + days_back, + enqueued + ); + Ok(enqueued) +} + +fn next_sleep_duration() -> Duration { + let now = Utc::now(); + let tomorrow = now.date_naive() + ChronoDuration::days(1); + let next = Utc + .with_ymd_and_hms(tomorrow.year(), tomorrow.month(), tomorrow.day(), 0, 5, 0) + // UTC has no DST gaps/overlaps, so `single()` always returns + // `Some` for any valid (Y, M, D, h, m, s). Fallback retained + // only as a defensive belt-and-braces against future API churn. + .single() + .unwrap_or_else(|| now + ChronoDuration::hours(24)); + (next - now) + .to_std() + .unwrap_or_else(|_| Duration::from_secs(60)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::jobs::store::{count_by_status, count_total, mark_done}; + use crate::openhuman::memory::tree::jobs::types::JobStatus; + use tempfile::TempDir; + + fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + (tmp, cfg) + } + + #[test] + fn trigger_digest_enqueues_a_job() { + let (_tmp, cfg) = test_config(); + let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap(); + let id = trigger_digest(&cfg, date).unwrap(); + assert!(id.is_some(), "first trigger must enqueue"); + assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1); + } + + #[test] + fn trigger_digest_dedupes_active_jobs() { + let (_tmp, cfg) = test_config(); + let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap(); + let first = trigger_digest(&cfg, date).unwrap(); + let second = trigger_digest(&cfg, date).unwrap(); + assert!(first.is_some()); + assert!( + second.is_none(), + "duplicate trigger must be dedupe-suppressed while active" + ); + assert_eq!(count_total(&cfg).unwrap(), 1); + } + + #[test] + fn trigger_digest_after_done_creates_fresh_row() { + let (_tmp, cfg) = test_config(); + let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap(); + let id1 = trigger_digest(&cfg, date).unwrap().unwrap(); + // Simulate a worker finishing the job — releases the dedupe slot. + mark_done(&cfg, &id1).unwrap(); + + let id2 = trigger_digest(&cfg, date).unwrap(); + assert!( + id2.is_some(), + "after the prior job completes, a fresh trigger must enqueue" + ); + assert_ne!(id2.unwrap(), id1); + assert_eq!(count_total(&cfg).unwrap(), 2); + } + + #[test] + fn backfill_missing_digests_enqueues_one_per_day() { + let (_tmp, cfg) = test_config(); + let n = backfill_missing_digests(&cfg, 5).unwrap(); + assert_eq!(n, 5, "expected one job per day in the 5-day window"); + assert_eq!(count_total(&cfg).unwrap(), 5); + } + + #[test] + fn backfill_missing_digests_zero_window_is_noop() { + let (_tmp, cfg) = test_config(); + let n = backfill_missing_digests(&cfg, 0).unwrap(); + assert_eq!(n, 0); + assert_eq!(count_total(&cfg).unwrap(), 0); + } + + #[test] + fn backfill_missing_digests_is_idempotent_while_active() { + let (_tmp, cfg) = test_config(); + let n1 = backfill_missing_digests(&cfg, 3).unwrap(); + let n2 = backfill_missing_digests(&cfg, 3).unwrap(); + assert_eq!(n1, 3); + assert_eq!(n2, 0, "second call must be fully dedupe-suppressed"); + assert_eq!(count_total(&cfg).unwrap(), 3); + } +} diff --git a/src/openhuman/memory/tree/jobs/store.rs b/src/openhuman/memory/tree/jobs/store.rs new file mode 100644 index 00000000000..1347e2a2de5 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/store.rs @@ -0,0 +1,486 @@ +//! SQLite persistence for the memory-tree job queue. +//! +//! Producers call [`enqueue`] inside their own writes (or with a fresh tx) +//! to atomically commit the side-effect plus its follow-up job. The worker +//! pool calls [`claim_next`] to lease a job, [`mark_done`] / [`mark_failed`] +//! to settle it, and [`recover_stale_locks`] on startup to flip rows whose +//! `locked_until_ms` expired without a settle. +//! +//! Concurrency: +//! - The dedupe key is enforced by a partial `UNIQUE` index that only +//! covers `status IN ('ready', 'running')`. Producers use `INSERT OR +//! IGNORE` so a duplicate enqueue while a job is in flight or queued is +//! a silent no-op; a duplicate enqueue after the first completes is +//! accepted and creates a fresh row. +//! - `claim_next` is one statement: `UPDATE … WHERE id = (SELECT … LIMIT 1) +//! RETURNING …`. SQLite serialises writes, so no two workers can claim +//! the same row. + +use anyhow::{Context, Result}; +use chrono::Utc; +use rusqlite::{params, Connection, OptionalExtension, Transaction}; +use uuid::Uuid; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::jobs::types::{Job, JobKind, JobStatus, NewJob}; +use crate::openhuman::memory::tree::store::with_connection; + +/// Default visibility lock — a worker that crashes mid-job will have its +/// row recovered after this window. 5 min is comfortably larger than any +/// expected single-job runtime (LLM extract or summarise) without leaving +/// real failures stuck for hours. +pub const DEFAULT_LOCK_DURATION_MS: i64 = 5 * 60 * 1_000; + +/// Backoff math for retry. Returns `now + min(base * 2^attempts, cap)`. +const RETRY_BASE_MS: i64 = 60 * 1_000; +const RETRY_CAP_MS: i64 = 60 * 60 * 1_000; +const DEFAULT_MAX_ATTEMPTS: u32 = 5; + +/// Enqueue one job. Idempotent on `dedupe_key` while another active row +/// (status `ready`/`running`) shares it. Returns `Some(id)` if the row +/// was inserted, `None` if a duplicate was suppressed. +pub fn enqueue(config: &Config, job: &NewJob) -> Result> { + with_connection(config, |conn| enqueue_conn(conn, job)) +} + +/// Enqueue inside a caller-owned transaction. Use this when the producer +/// is already mid-tx (e.g. `ingest::persist` writing chunks + jobs in one +/// commit) so the queue insert lands atomically with the side-effect. +/// `Transaction` derefs to `Connection`, so callers just pass `&tx`. +pub fn enqueue_tx(tx: &Transaction<'_>, job: &NewJob) -> Result> { + enqueue_conn(tx, job) +} + +pub(crate) fn enqueue_conn(conn: &Connection, job: &NewJob) -> Result> { + let id = format!("job:{}", Uuid::new_v4()); + let now_ms = Utc::now().timestamp_millis(); + let available_at = job.available_at_ms.unwrap_or(now_ms); + let max_attempts = job.max_attempts.unwrap_or(DEFAULT_MAX_ATTEMPTS) as i64; + + let inserted = conn.execute( + "INSERT OR IGNORE INTO mem_tree_jobs ( + id, kind, payload_json, dedupe_key, status, attempts, max_attempts, + available_at_ms, locked_until_ms, last_error, + created_at_ms, started_at_ms, completed_at_ms + ) VALUES (?1, ?2, ?3, ?4, 'ready', 0, ?5, ?6, NULL, NULL, ?7, NULL, NULL)", + params![ + id, + job.kind.as_str(), + job.payload_json, + job.dedupe_key, + max_attempts, + available_at, + now_ms, + ], + )?; + + if inserted == 0 { + log::debug!( + "[memory_tree::jobs] enqueue suppressed by dedupe kind={} key={:?}", + job.kind.as_str(), + job.dedupe_key + ); + return Ok(None); + } + log::debug!( + "[memory_tree::jobs] enqueued id={} kind={} avail_at_ms={} dedupe={:?}", + id, + job.kind.as_str(), + available_at, + job.dedupe_key + ); + Ok(Some(id)) +} + +/// Atomically claim the next ready job whose `available_at_ms` has come +/// due. Sets `status=running`, bumps `attempts`, stamps `started_at_ms` +/// and `locked_until_ms`. Returns `None` when the queue is empty / not +/// yet due. +pub fn claim_next(config: &Config, lock_duration_ms: i64) -> Result> { + with_connection(config, |conn| { + let now_ms = Utc::now().timestamp_millis(); + let lock_until = now_ms.saturating_add(lock_duration_ms); + + let row = conn + .query_row( + "UPDATE mem_tree_jobs + SET status = 'running', + attempts = attempts + 1, + started_at_ms = ?1, + locked_until_ms = ?2, + last_error = NULL + WHERE id = ( + SELECT id FROM mem_tree_jobs + WHERE status = 'ready' + AND available_at_ms <= ?1 + ORDER BY available_at_ms ASC + LIMIT 1 + ) + RETURNING id, kind, payload_json, dedupe_key, status, attempts, + max_attempts, available_at_ms, locked_until_ms, last_error, + created_at_ms, started_at_ms, completed_at_ms", + params![now_ms, lock_until], + row_to_job, + ) + .optional() + .context("Failed to claim next mem_tree_jobs row")?; + if let Some(j) = &row { + log::debug!( + "[memory_tree::jobs] claimed id={} kind={} attempt={}/{}", + j.id, + j.kind.as_str(), + j.attempts, + j.max_attempts + ); + } + Ok(row) + }) +} + +/// Mark a claimed job as `done`. Clears the lock and stamps `completed_at_ms`. +pub fn mark_done(config: &Config, job_id: &str) -> Result<()> { + with_connection(config, |conn| { + let now_ms = Utc::now().timestamp_millis(); + let n = conn.execute( + "UPDATE mem_tree_jobs + SET status = 'done', + completed_at_ms = ?1, + locked_until_ms = NULL, + last_error = NULL + WHERE id = ?2", + params![now_ms, job_id], + )?; + if n == 0 { + log::warn!("[memory_tree::jobs] mark_done id={job_id} affected 0 rows"); + } + Ok(()) + }) +} + +/// Settle a failed job. If `attempts < max_attempts`, the row goes back +/// to `ready` with an exponential-backoff `available_at_ms`. Otherwise +/// it terminates as `failed`. Either way `last_error` is recorded. +pub fn mark_failed(config: &Config, job_id: &str, error: &str) -> Result<()> { + with_connection(config, |conn| { + let now_ms = Utc::now().timestamp_millis(); + let row: Option<(i64, i64)> = conn + .query_row( + "SELECT attempts, max_attempts FROM mem_tree_jobs WHERE id = ?1", + params![job_id], + |r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?)), + ) + .optional() + .context("Failed to read job for mark_failed")?; + + let Some((attempts, max_attempts)) = row else { + log::warn!("[memory_tree::jobs] mark_failed id={job_id} not found"); + return Ok(()); + }; + + if attempts >= max_attempts { + log::warn!( + "[memory_tree::jobs] terminal failure id={job_id} attempts={attempts}/{max_attempts} err={error}" + ); + conn.execute( + "UPDATE mem_tree_jobs + SET status = 'failed', + completed_at_ms = ?1, + locked_until_ms = NULL, + last_error = ?2 + WHERE id = ?3", + params![now_ms, error, job_id], + )?; + } else { + let backoff = backoff_ms(attempts as u32); + let next_at = now_ms.saturating_add(backoff); + log::info!( + "[memory_tree::jobs] retry id={job_id} attempt={attempts}/{max_attempts} \ + next_at_ms={next_at} err={error}" + ); + conn.execute( + "UPDATE mem_tree_jobs + SET status = 'ready', + available_at_ms = ?1, + locked_until_ms = NULL, + last_error = ?2 + WHERE id = ?3", + params![next_at, error, job_id], + )?; + } + Ok(()) + }) +} + +/// Flip any `running` row whose `locked_until_ms` has expired back to +/// `ready`. Called once at worker startup so a process crash mid-job +/// doesn't leave work stranded. Returns the number of rows recovered. +pub fn recover_stale_locks(config: &Config) -> Result { + with_connection(config, |conn| { + let now_ms = Utc::now().timestamp_millis(); + let n = conn.execute( + "UPDATE mem_tree_jobs + SET status = 'ready', + last_error = COALESCE(last_error, 'recovered_from_stale_lock') + WHERE status = 'running' + AND locked_until_ms IS NOT NULL + AND locked_until_ms < ?1", + params![now_ms], + )?; + if n > 0 { + log::warn!("[memory_tree::jobs] recovered {n} stale-locked job(s) at startup"); + } + Ok(n) + }) +} + +/// Quick count helper for tests / diagnostics. +pub fn count_by_status(config: &Config, status: JobStatus) -> Result { + with_connection(config, |conn| { + let n: i64 = conn.query_row( + "SELECT COUNT(*) FROM mem_tree_jobs WHERE status = ?1", + params![status.as_str()], + |r| r.get(0), + )?; + Ok(n.max(0) as u64) + }) +} + +/// Total count regardless of status — handy for assertions. +pub fn count_total(config: &Config) -> Result { + with_connection(config, |conn| { + let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| r.get(0))?; + Ok(n.max(0) as u64) + }) +} + +/// Fetch one job by id (test/diagnostic helper). +pub fn get_job(config: &Config, id: &str) -> Result> { + with_connection(config, |conn| { + let job = conn + .query_row( + "SELECT id, kind, payload_json, dedupe_key, status, attempts, max_attempts, + available_at_ms, locked_until_ms, last_error, + created_at_ms, started_at_ms, completed_at_ms + FROM mem_tree_jobs WHERE id = ?1", + params![id], + row_to_job, + ) + .optional()?; + Ok(job) + }) +} + +fn row_to_job(row: &rusqlite::Row<'_>) -> rusqlite::Result { + let id: String = row.get(0)?; + let kind_s: String = row.get(1)?; + let payload_json: String = row.get(2)?; + let dedupe_key: Option = row.get(3)?; + let status_s: String = row.get(4)?; + let attempts: i64 = row.get(5)?; + let max_attempts: i64 = row.get(6)?; + let available_at_ms: i64 = row.get(7)?; + let locked_until_ms: Option = row.get(8)?; + let last_error: Option = row.get(9)?; + let created_at_ms: i64 = row.get(10)?; + let started_at_ms: Option = row.get(11)?; + let completed_at_ms: Option = row.get(12)?; + + let kind = JobKind::parse(&kind_s).map_err(|e| { + rusqlite::Error::FromSqlConversionFailure(1, rusqlite::types::Type::Text, e.into()) + })?; + let status = JobStatus::parse(&status_s).map_err(|e| { + rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, e.into()) + })?; + + Ok(Job { + id, + kind, + payload_json, + dedupe_key, + status, + attempts: attempts.max(0) as u32, + max_attempts: max_attempts.max(0) as u32, + available_at_ms, + locked_until_ms, + last_error, + created_at_ms, + started_at_ms, + completed_at_ms, + }) +} + +/// Exponential backoff: attempt 1 → 60s, 2 → 120s, 3 → 240s, capped at 1h. +fn backoff_ms(attempts_so_far: u32) -> i64 { + // attempts_so_far is the count BEFORE the next retry's attempt — so the + // first retry uses attempts_so_far=1, giving base*2^0 = 60s. + let exp = attempts_so_far.saturating_sub(1).min(20); // cap shift + let mult = 1i64 << exp; // 1, 2, 4, … + let raw = RETRY_BASE_MS.saturating_mul(mult); + raw.min(RETRY_CAP_MS) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::jobs::types::{ + AppendBufferPayload, AppendTarget, ExtractChunkPayload, NodeRef, + }; + use tempfile::TempDir; + + fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + (tmp, cfg) + } + + #[test] + fn enqueue_and_claim_roundtrip() { + let (_tmp, cfg) = test_config(); + let payload = ExtractChunkPayload { + chunk_id: "c1".into(), + }; + let nj = NewJob::extract_chunk(&payload).unwrap(); + let id = enqueue(&cfg, &nj).unwrap().expect("inserted"); + + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + assert_eq!(claimed.id, id); + assert_eq!(claimed.status, JobStatus::Running); + assert_eq!(claimed.attempts, 1); + assert!(claimed.locked_until_ms.is_some()); + + // Second claim should see no eligible row (the only one is now running). + let again = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap(); + assert!(again.is_none()); + + mark_done(&cfg, &id).unwrap(); + let row = get_job(&cfg, &id).unwrap().unwrap(); + assert_eq!(row.status, JobStatus::Done); + assert!(row.completed_at_ms.is_some()); + assert!(row.locked_until_ms.is_none()); + } + + #[test] + fn enqueue_dedupes_active_jobs() { + let (_tmp, cfg) = test_config(); + let payload = ExtractChunkPayload { + chunk_id: "c1".into(), + }; + let nj = NewJob::extract_chunk(&payload).unwrap(); + let id1 = enqueue(&cfg, &nj).unwrap(); + let id2 = enqueue(&cfg, &nj).unwrap(); + assert!(id1.is_some()); + assert!(id2.is_none(), "duplicate should be suppressed while ready"); + assert_eq!(count_total(&cfg).unwrap(), 1); + } + + #[test] + fn enqueue_after_done_creates_fresh_row() { + let (_tmp, cfg) = test_config(); + let payload = ExtractChunkPayload { + chunk_id: "c1".into(), + }; + let nj = NewJob::extract_chunk(&payload).unwrap(); + let id1 = enqueue(&cfg, &nj).unwrap().unwrap(); + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + assert_eq!(claimed.id, id1); + mark_done(&cfg, &id1).unwrap(); + + // Now the dedupe key is free (partial index excludes 'done'). + let id2 = enqueue(&cfg, &nj).unwrap(); + assert!(id2.is_some()); + assert_ne!(id2.unwrap(), id1); + assert_eq!(count_total(&cfg).unwrap(), 2); + } + + #[test] + fn mark_failed_retries_then_terminates() { + let (_tmp, cfg) = test_config(); + let payload = AppendBufferPayload { + node: NodeRef::Leaf { + chunk_id: "c1".into(), + }, + target: AppendTarget::Source { + source_id: "slack:#x".into(), + }, + }; + let mut nj = NewJob::append_buffer(&payload).unwrap(); + nj.max_attempts = Some(2); + let id = enqueue(&cfg, &nj).unwrap().unwrap(); + + // Fail #1 — should bounce back to 'ready' with future available_at. + let _ = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + mark_failed(&cfg, &id, "boom").unwrap(); + let row = get_job(&cfg, &id).unwrap().unwrap(); + assert_eq!(row.status, JobStatus::Ready); + assert!(row.available_at_ms > Utc::now().timestamp_millis()); + assert_eq!(row.last_error.as_deref(), Some("boom")); + + // Force the row available again so the test doesn't hinge on sleep. + with_connection(&cfg, |c| { + c.execute( + "UPDATE mem_tree_jobs SET available_at_ms = 0 WHERE id = ?1", + params![id], + )?; + Ok(()) + }) + .unwrap(); + + // Fail #2 — exceeds max_attempts → terminal 'failed'. + let _ = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + mark_failed(&cfg, &id, "fatal").unwrap(); + let row = get_job(&cfg, &id).unwrap().unwrap(); + assert_eq!(row.status, JobStatus::Failed); + assert_eq!(row.last_error.as_deref(), Some("fatal")); + assert!(row.completed_at_ms.is_some()); + } + + #[test] + fn recover_stale_locks_resets_running_rows() { + let (_tmp, cfg) = test_config(); + let payload = ExtractChunkPayload { + chunk_id: "c1".into(), + }; + let nj = NewJob::extract_chunk(&payload).unwrap(); + let id = enqueue(&cfg, &nj).unwrap().unwrap(); + + // Claim with a lock window that's already in the past so recovery + // sees it as expired. + let _ = claim_next(&cfg, -1).unwrap().unwrap(); + + let recovered = recover_stale_locks(&cfg).unwrap(); + assert_eq!(recovered, 1); + let row = get_job(&cfg, &id).unwrap().unwrap(); + assert_eq!(row.status, JobStatus::Ready); + } + + #[test] + fn backoff_grows_then_caps() { + assert_eq!(backoff_ms(1), 60_000); + assert_eq!(backoff_ms(2), 120_000); + assert_eq!(backoff_ms(3), 240_000); + // Eventually clamps at the cap. + assert_eq!(backoff_ms(20), RETRY_CAP_MS); + assert_eq!(backoff_ms(99), RETRY_CAP_MS); + } + + #[test] + fn count_by_status_reports_each_state() { + let (_tmp, cfg) = test_config(); + for i in 0..3 { + let p = ExtractChunkPayload { + chunk_id: format!("c{i}"), + }; + let nj = NewJob::extract_chunk(&p).unwrap(); + enqueue(&cfg, &nj).unwrap(); + } + assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 3); + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + mark_done(&cfg, &claimed.id).unwrap(); + assert_eq!(count_by_status(&cfg, JobStatus::Done).unwrap(), 1); + assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 2); + } +} diff --git a/src/openhuman/memory/tree/jobs/testing.rs b/src/openhuman/memory/tree/jobs/testing.rs new file mode 100644 index 00000000000..e8211fb7893 --- /dev/null +++ b/src/openhuman/memory/tree/jobs/testing.rs @@ -0,0 +1,15 @@ +use anyhow::Result; + +use crate::openhuman::config::Config; + +/// Deterministically run queued memory-tree jobs until no immediately +/// claimable work remains. Intended for tests that need the async pipeline +/// to settle without spawning background tasks. +pub async fn drain_until_idle(config: &Config) -> Result<()> { + loop { + if !super::worker::run_once(config).await? { + break; + } + } + Ok(()) +} diff --git a/src/openhuman/memory/tree/jobs/types.rs b/src/openhuman/memory/tree/jobs/types.rs new file mode 100644 index 00000000000..e6bf044f21c --- /dev/null +++ b/src/openhuman/memory/tree/jobs/types.rs @@ -0,0 +1,439 @@ +//! Job types for the async memory-tree pipeline. +//! +//! Each `Job` row in `mem_tree_jobs` stores its discriminator as a string +//! `kind` plus a JSON-encoded `payload`. The strongly-typed payload structs +//! below own (de)serialisation; handlers parse the payload by branching on +//! [`JobKind`] and calling the matching `from_payload_json`. + +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; + +/// Discriminator persisted in `mem_tree_jobs.kind`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum JobKind { + /// Run LLM entity extraction over a single chunk and decide admission. + ExtractChunk, + /// Push an admitted chunk into a tree's L0 buffer. + AppendBuffer, + /// Seal exactly one buffer level; cascades enqueue a follow-up. + Seal, + /// Match a chunk's entities against active topic trees and enqueue + /// per-topic `AppendBuffer` jobs. + TopicRoute, + /// Build the global tree's daily digest for a given UTC date. + DigestDaily, + /// Walk stale buffers and enqueue `Seal` jobs for any over the age cap. + FlushStale, +} + +impl JobKind { + pub fn as_str(&self) -> &'static str { + match self { + JobKind::ExtractChunk => "extract_chunk", + JobKind::AppendBuffer => "append_buffer", + JobKind::Seal => "seal", + JobKind::TopicRoute => "topic_route", + JobKind::DigestDaily => "digest_daily", + JobKind::FlushStale => "flush_stale", + } + } + + pub fn parse(s: &str) -> Result { + Ok(match s { + "extract_chunk" => JobKind::ExtractChunk, + "append_buffer" => JobKind::AppendBuffer, + "seal" => JobKind::Seal, + "topic_route" => JobKind::TopicRoute, + "digest_daily" => JobKind::DigestDaily, + "flush_stale" => JobKind::FlushStale, + other => return Err(anyhow!("unknown JobKind '{other}'")), + }) + } + + /// True when handling this kind should hold a slot from the global + /// LLM concurrency semaphore. `TopicRoute` is bound because + /// `maybe_spawn_topic_tree → backfill_topic_tree` can transitively + /// trigger summariser LLM calls when an entity first crosses the + /// hotness threshold. + pub fn is_llm_bound(&self) -> bool { + matches!( + self, + JobKind::ExtractChunk | JobKind::Seal | JobKind::DigestDaily | JobKind::TopicRoute + ) + } +} + +/// Lifecycle states persisted on `mem_tree_jobs.status`. Workers transition +/// `ready → running → done|failed`. `Cancelled` is reserved for explicit +/// admin actions (none surfaced yet). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JobStatus { + Ready, + Running, + Done, + Failed, + Cancelled, +} + +impl JobStatus { + pub fn as_str(&self) -> &'static str { + match self { + JobStatus::Ready => "ready", + JobStatus::Running => "running", + JobStatus::Done => "done", + JobStatus::Failed => "failed", + JobStatus::Cancelled => "cancelled", + } + } + + pub fn parse(s: &str) -> Result { + Ok(match s { + "ready" => JobStatus::Ready, + "running" => JobStatus::Running, + "done" => JobStatus::Done, + "failed" => JobStatus::Failed, + "cancelled" => JobStatus::Cancelled, + other => return Err(anyhow!("unknown JobStatus '{other}'")), + }) + } + + pub fn is_terminal(&self) -> bool { + matches!( + self, + JobStatus::Done | JobStatus::Failed | JobStatus::Cancelled + ) + } +} + +// ── Payloads ─────────────────────────────────────────────────────────────── + +/// Reference to either a leaf chunk or a sealed summary node. Used by +/// payloads that route content through the pipeline regardless of which +/// kind of source produced it. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum NodeRef { + Leaf { chunk_id: String }, + Summary { summary_id: String }, +} + +impl NodeRef { + /// Stringified id with kind prefix, suitable for dedupe-key composition. + pub fn dedupe_fragment(&self) -> String { + match self { + NodeRef::Leaf { chunk_id } => format!("leaf:{chunk_id}"), + NodeRef::Summary { summary_id } => format!("summary:{summary_id}"), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ExtractChunkPayload { + pub chunk_id: String, +} + +impl ExtractChunkPayload { + pub fn dedupe_key(&self) -> String { + format!("extract:{}", self.chunk_id) + } +} + +/// Where an `AppendBuffer` job should land its node. Source-tree appends +/// are keyed by `source_id`; topic-tree appends are keyed by `tree_id` +/// because there can be many topic trees per node. +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum AppendTarget { + Source { source_id: String }, + Topic { tree_id: String }, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct AppendBufferPayload { + pub node: NodeRef, + pub target: AppendTarget, +} + +impl AppendBufferPayload { + pub fn dedupe_key(&self) -> String { + let node_part = self.node.dedupe_fragment(); + match &self.target { + AppendTarget::Source { source_id } => { + format!("append:source:{source_id}:{node_part}") + } + AppendTarget::Topic { tree_id } => { + format!("append:topic:{tree_id}:{node_part}") + } + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SealPayload { + pub tree_id: String, + pub level: u32, + /// When `Some`, the seal handler bypasses the buffer-budget check and + /// force-seals — used by the time-based flush path. The wall-clock is + /// passed through so the seal stamps a deterministic `sealed_at`. + pub force_now_ms: Option, +} + +impl SealPayload { + pub fn dedupe_key(&self) -> String { + // Active seal-job uniqueness is enforced per (tree, level): a seal + // already in flight suppresses duplicate enqueues. Once the job + // completes the partial index releases the key, so the next time + // the buffer crosses its gate a fresh seal can be enqueued. + format!("seal:{}:{}", self.tree_id, self.level) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TopicRoutePayload { + pub node: NodeRef, +} + +impl TopicRoutePayload { + pub fn dedupe_key(&self) -> String { + format!("topic_route:{}", self.node.dedupe_fragment()) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct DigestDailyPayload { + /// UTC calendar date in `YYYY-MM-DD` form. Stored as a string so the + /// dedupe key doesn't need to know about chrono. + pub date_iso: String, +} + +impl DigestDailyPayload { + pub fn dedupe_key(&self) -> String { + format!("digest_daily:{}", self.date_iso) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, Default)] +pub struct FlushStalePayload { + /// Override the configured `DEFAULT_FLUSH_AGE_SECS`. Optional so the + /// scheduler can enqueue with `None` and let the handler use the + /// configured default. + pub max_age_secs: Option, +} + +impl FlushStalePayload { + pub fn dedupe_key(&self, date_iso: &str) -> String { + format!("flush_stale:{date_iso}") + } +} + +/// One row in `mem_tree_jobs`. `payload_json` is left as a raw string so +/// callers parse it lazily based on `kind`. +#[derive(Clone, Debug)] +pub struct Job { + pub id: String, + pub kind: JobKind, + pub payload_json: String, + pub dedupe_key: Option, + pub status: JobStatus, + pub attempts: u32, + pub max_attempts: u32, + pub available_at_ms: i64, + pub locked_until_ms: Option, + pub last_error: Option, + pub created_at_ms: i64, + pub started_at_ms: Option, + pub completed_at_ms: Option, +} + +/// Caller-side bundle for `enqueue` — `Job` minus the persistence-only +/// columns. Keeps producers from having to mint timestamps and ids by hand. +#[derive(Clone, Debug)] +pub struct NewJob { + pub kind: JobKind, + pub payload_json: String, + pub dedupe_key: Option, + /// `None` means "available immediately." Set this for delayed jobs + /// (retries, scheduled work). + pub available_at_ms: Option, + pub max_attempts: Option, +} + +impl NewJob { + pub fn extract_chunk(p: &ExtractChunkPayload) -> Result { + Ok(Self { + kind: JobKind::ExtractChunk, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key()), + available_at_ms: None, + max_attempts: None, + }) + } + + pub fn append_buffer(p: &AppendBufferPayload) -> Result { + Ok(Self { + kind: JobKind::AppendBuffer, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key()), + available_at_ms: None, + max_attempts: None, + }) + } + + pub fn seal(p: &SealPayload) -> Result { + Ok(Self { + kind: JobKind::Seal, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key()), + available_at_ms: None, + max_attempts: None, + }) + } + + pub fn topic_route(p: &TopicRoutePayload) -> Result { + Ok(Self { + kind: JobKind::TopicRoute, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key()), + available_at_ms: None, + max_attempts: None, + }) + } + + pub fn digest_daily(p: &DigestDailyPayload) -> Result { + Ok(Self { + kind: JobKind::DigestDaily, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key()), + available_at_ms: None, + max_attempts: None, + }) + } + + pub fn flush_stale(p: &FlushStalePayload, date_iso: &str) -> Result { + Ok(Self { + kind: JobKind::FlushStale, + payload_json: serde_json::to_string(p)?, + dedupe_key: Some(p.dedupe_key(date_iso)), + available_at_ms: None, + max_attempts: None, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn job_kind_roundtrip() { + for k in [ + JobKind::ExtractChunk, + JobKind::AppendBuffer, + JobKind::Seal, + JobKind::TopicRoute, + JobKind::DigestDaily, + JobKind::FlushStale, + ] { + assert_eq!(JobKind::parse(k.as_str()).unwrap(), k); + } + } + + #[test] + fn job_status_terminality() { + assert!(!JobStatus::Ready.is_terminal()); + assert!(!JobStatus::Running.is_terminal()); + assert!(JobStatus::Done.is_terminal()); + assert!(JobStatus::Failed.is_terminal()); + assert!(JobStatus::Cancelled.is_terminal()); + } + + #[test] + fn dedupe_keys_distinguish_targets() { + let p_src = AppendBufferPayload { + node: NodeRef::Leaf { + chunk_id: "c1".into(), + }, + target: AppendTarget::Source { + source_id: "slack:#eng".into(), + }, + }; + let p_topic = AppendBufferPayload { + node: NodeRef::Leaf { + chunk_id: "c1".into(), + }, + target: AppendTarget::Topic { + tree_id: "topic:abc".into(), + }, + }; + assert_ne!(p_src.dedupe_key(), p_topic.dedupe_key()); + } + + #[test] + fn dedupe_keys_distinguish_node_kinds() { + let p_leaf = AppendBufferPayload { + node: NodeRef::Leaf { + chunk_id: "x".into(), + }, + target: AppendTarget::Topic { + tree_id: "t".into(), + }, + }; + let p_summary = AppendBufferPayload { + node: NodeRef::Summary { + summary_id: "x".into(), + }, + target: AppendTarget::Topic { + tree_id: "t".into(), + }, + }; + assert_ne!(p_leaf.dedupe_key(), p_summary.dedupe_key()); + + let r_leaf = TopicRoutePayload { + node: NodeRef::Leaf { + chunk_id: "x".into(), + }, + }; + let r_summary = TopicRoutePayload { + node: NodeRef::Summary { + summary_id: "x".into(), + }, + }; + assert_ne!(r_leaf.dedupe_key(), r_summary.dedupe_key()); + } + + #[test] + fn llm_bound_kinds() { + assert!(JobKind::ExtractChunk.is_llm_bound()); + assert!(JobKind::Seal.is_llm_bound()); + assert!(JobKind::DigestDaily.is_llm_bound()); + assert!(JobKind::TopicRoute.is_llm_bound()); + assert!(!JobKind::AppendBuffer.is_llm_bound()); + assert!(!JobKind::FlushStale.is_llm_bound()); + } + + #[test] + fn node_ref_serializes_with_kind_tag() { + let leaf = NodeRef::Leaf { + chunk_id: "x".into(), + }; + let s = serde_json::to_string(&leaf).unwrap(); + assert!(s.contains("\"kind\":\"leaf\"")); + let back: NodeRef = serde_json::from_str(&s).unwrap(); + assert_eq!(back, leaf); + } + + #[test] + fn append_target_serializes_with_kind_tag() { + let p = AppendTarget::Source { + source_id: "x".into(), + }; + let s = serde_json::to_string(&p).unwrap(); + assert!(s.contains("\"kind\":\"source\"")); + assert!(s.contains("\"source_id\":\"x\"")); + let back: AppendTarget = serde_json::from_str(&s).unwrap(); + match back { + AppendTarget::Source { source_id } => assert_eq!(source_id, "x"), + _ => panic!("wrong variant"), + } + } +} diff --git a/src/openhuman/memory/tree/jobs/worker.rs b/src/openhuman/memory/tree/jobs/worker.rs new file mode 100644 index 00000000000..29391c5057f --- /dev/null +++ b/src/openhuman/memory/tree/jobs/worker.rs @@ -0,0 +1,109 @@ +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use anyhow::Result; +use tokio::sync::{Notify, Semaphore}; + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::jobs::handlers; +use crate::openhuman::memory::tree::jobs::store::{ + claim_next, mark_done, mark_failed, recover_stale_locks, DEFAULT_LOCK_DURATION_MS, +}; + +const WORKER_COUNT: usize = 3; +const POLL_INTERVAL: Duration = Duration::from_secs(5); + +static WORKER_NOTIFY: OnceLock> = OnceLock::new(); +static STARTED: std::sync::Once = std::sync::Once::new(); + +pub fn wake_workers() { + if let Some(notify) = WORKER_NOTIFY.get() { + notify.notify_waiters(); + } +} + +/// Start the worker pool + daily scheduler. Takes the full `Config` so +/// each spawned task sees the user's actual settings (LLM endpoints, +/// embedder model, timeouts) — not `Config::default()`. Without this, +/// workers fall back to inert/regex-only behavior regardless of what's +/// in `config.toml`, defeating the entire async pipeline. +/// +/// Idempotent (`Once`-guarded) so repeat calls during bootstrap are +/// safe no-ops after the first. +pub fn start(config: Config) { + STARTED.call_once(|| { + let notify = WORKER_NOTIFY + .get_or_init(|| Arc::new(Notify::new())) + .clone(); + let llm_slots = Arc::new(Semaphore::new(3)); + if let Err(err) = recover_stale_locks(&config) { + log::warn!("[memory_tree::jobs] recover_stale_locks failed at startup: {err:#}"); + } + + for idx in 0..WORKER_COUNT { + let notify = notify.clone(); + let llm_slots = llm_slots.clone(); + let cfg = config.clone(); + tokio::spawn(async move { + loop { + match run_once_with_semaphore(&cfg, llm_slots.clone()).await { + Ok(true) => continue, + Ok(false) => { + tokio::select! { + _ = notify.notified() => {} + _ = tokio::time::sleep(POLL_INTERVAL) => {} + } + } + Err(err) => { + log::error!("[memory_tree::jobs] worker={} loop error: {:#}", idx, err); + tokio::time::sleep(Duration::from_secs(1)).await; + } + } + } + }); + } + + super::scheduler::start(config); + }); +} + +pub async fn run_once(config: &Config) -> Result { + let llm_slots = Arc::new(Semaphore::new(1)); + run_once_with_semaphore(config, llm_slots).await +} + +async fn run_once_with_semaphore(config: &Config, llm_slots: Arc) -> Result { + let Some(job) = claim_next(config, DEFAULT_LOCK_DURATION_MS)? else { + return Ok(false); + }; + + let permit = if job.kind.is_llm_bound() { + Some(llm_slots.acquire().await?) + } else { + None + }; + let result = handlers::handle_job(config, &job).await; + drop(permit); + + match result { + Ok(()) => { + mark_done(config, &job.id)?; + } + Err(err) => { + // Preserve the full anyhow cause chain in the persisted + // last_error so a reader of mem_tree_jobs can see the root + // cause, not just the top-level message. Mirrors the {:#} + // log format used right above. + let message = format!("{err:#}"); + log::warn!( + "[memory_tree::jobs] job failed id={} kind={} err={:#}", + job.id, + job.kind.as_str(), + err + ); + mark_failed(config, &job.id, &message)?; + } + } + + Ok(true) +} diff --git a/src/openhuman/memory/tree/mod.rs b/src/openhuman/memory/tree/mod.rs index 038b352d31e..7538050882d 100644 --- a/src/openhuman/memory/tree/mod.rs +++ b/src/openhuman/memory/tree/mod.rs @@ -26,6 +26,7 @@ pub mod canonicalize; pub mod chunker; pub mod global_tree; pub mod ingest; +pub mod jobs; pub mod retrieval; pub mod rpc; pub mod schemas; diff --git a/src/openhuman/memory/tree/retrieval/drill_down.rs b/src/openhuman/memory/tree/retrieval/drill_down.rs index 9a307bfb149..5243abf4706 100644 --- a/src/openhuman/memory/tree/retrieval/drill_down.rs +++ b/src/openhuman/memory/tree/retrieval/drill_down.rs @@ -215,7 +215,9 @@ fn walk_with_embeddings( #[cfg(test)] mod tests { use super::*; - use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, + }; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; use crate::openhuman::memory::tree::source_tree::types::TreeKind; @@ -273,6 +275,7 @@ mod tests { score: 0.5, }, &summariser, + &LabelStrategy::Empty, ) .await .unwrap(); diff --git a/src/openhuman/memory/tree/retrieval/global.rs b/src/openhuman/memory/tree/retrieval/global.rs index fd5199f2944..f52a3a9a920 100644 --- a/src/openhuman/memory/tree/retrieval/global.rs +++ b/src/openhuman/memory/tree/retrieval/global.rs @@ -88,7 +88,9 @@ fn recap_to_hits(recap: RecapOutput, tree_id: &str, tree_scope: &str) -> Vec= 1, "expected at least one kept chunk"); + assert!( + out.chunks_written >= 1, + "expected at least one persisted chunk" + ); + drain_until_idle(&cfg).await.unwrap(); for id in &out.chunk_ids { let emb = get_chunk_embedding(&cfg, id).unwrap(); let v = emb.unwrap_or_else(|| panic!("embedding missing for chunk_id={id}")); @@ -189,7 +193,9 @@ async fn ingest_populates_chunk_embeddings() { #[tokio::test] async fn seal_populates_summary_embedding() { use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM; - use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; + use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, + }; use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree; use crate::openhuman::memory::tree::source_tree::store as src_store; use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser; @@ -230,12 +236,24 @@ async fn seal_populates_summary_embedding() { topics: vec![], score: 0.5, }; - append_leaf(&cfg, &tree, &leaf_of(&c1), &summariser) - .await - .unwrap(); - let sealed = append_leaf(&cfg, &tree, &leaf_of(&c2), &summariser) - .await - .unwrap(); + append_leaf( + &cfg, + &tree, + &leaf_of(&c1), + &summariser, + &LabelStrategy::Empty, + ) + .await + .unwrap(); + let sealed = append_leaf( + &cfg, + &tree, + &leaf_of(&c2), + &summariser, + &LabelStrategy::Empty, + ) + .await + .unwrap(); assert_eq!(sealed.len(), 1, "expected one seal at the budget crossing"); let summary = src_store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); diff --git a/src/openhuman/memory/tree/retrieval/source.rs b/src/openhuman/memory/tree/retrieval/source.rs index 46d5bc4736b..3982e3568df 100644 --- a/src/openhuman/memory/tree/retrieval/source.rs +++ b/src/openhuman/memory/tree/retrieval/source.rs @@ -293,7 +293,9 @@ fn filter_by_window(hits: Vec, window_days: u32) -> Vec, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TriggerDigestResponse { + /// True when the job was newly enqueued; false when an active job for + /// the same date was suppressed by the dedupe partial unique index. + pub enqueued: bool, + /// ID of the freshly-inserted job row (None when dedupe-suppressed). + pub job_id: Option, + /// The actual date the digest will run for, echoed back as + /// `YYYY-MM-DD`. Useful when the caller didn't pass `date_iso` and + /// wants to know what default got chosen. + pub date_iso: String, +} + +pub async fn trigger_digest_rpc( + config: &Config, + req: TriggerDigestRequest, +) -> Result, String> { + use crate::openhuman::memory::tree::jobs; + use chrono::{Duration as ChronoDuration, NaiveDate, Utc}; + + let date = match req + .date_iso + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + Some(s) => NaiveDate::parse_from_str(s, "%Y-%m-%d") + .map_err(|e| format!("invalid date_iso (expected YYYY-MM-DD): {e}"))?, + None => Utc::now().date_naive() - ChronoDuration::days(1), + }; + let date_iso = date.format("%Y-%m-%d").to_string(); + + // Run the synchronous enqueue on a blocking thread — `trigger_digest` + // touches SQLite and we don't want to block the async runtime even + // for the few-microsecond INSERT. + let cfg_clone = config.clone(); + let date_for_blocking = date; + let job_id = + tokio::task::spawn_blocking(move || jobs::trigger_digest(&cfg_clone, date_for_blocking)) + .await + .map_err(|e| format!("trigger_digest join error: {e}"))? + .map_err(|e| format!("trigger_digest: {e}"))?; + + let enqueued = job_id.is_some(); + Ok(RpcOutcome::single_log( + TriggerDigestResponse { + enqueued, + job_id, + date_iso: date_iso.clone(), + }, + format!("memory_tree: trigger_digest date={date_iso} enqueued={enqueued}"), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::memory::tree::jobs::store::count_total; + use chrono::{Duration as ChronoDuration, Utc}; + use tempfile::TempDir; + + fn test_config() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + (tmp, cfg) + } + + #[tokio::test] + async fn trigger_digest_with_explicit_date_enqueues() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest { + date_iso: Some("2026-04-27".into()), + }; + let outcome = trigger_digest_rpc(&cfg, req).await.unwrap(); + let resp = outcome.value; + assert!(resp.enqueued); + assert!(resp.job_id.is_some()); + assert_eq!(resp.date_iso, "2026-04-27"); + assert_eq!(count_total(&cfg).unwrap(), 1); + } + + #[tokio::test] + async fn trigger_digest_with_no_date_defaults_to_yesterday() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest::default(); + let outcome = trigger_digest_rpc(&cfg, req).await.unwrap(); + let resp = outcome.value; + assert!(resp.enqueued); + let expected = (Utc::now().date_naive() - ChronoDuration::days(1)) + .format("%Y-%m-%d") + .to_string(); + assert_eq!(resp.date_iso, expected); + } + + #[tokio::test] + async fn trigger_digest_rejects_malformed_date() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest { + date_iso: Some("not-a-date".into()), + }; + let err = trigger_digest_rpc(&cfg, req).await.unwrap_err(); + assert!( + err.contains("invalid date_iso"), + "expected schema-shaped error message, got: {err}" + ); + assert_eq!(count_total(&cfg).unwrap(), 0); + } + + #[tokio::test] + async fn trigger_digest_dedupes_active_jobs() { + let (_tmp, cfg) = test_config(); + let req = TriggerDigestRequest { + date_iso: Some("2026-04-27".into()), + }; + let first = trigger_digest_rpc(&cfg, req.clone()).await.unwrap().value; + let second = trigger_digest_rpc(&cfg, req).await.unwrap().value; + assert!(first.enqueued); + assert!(!second.enqueued, "duplicate must be dedupe-suppressed"); + assert!(second.job_id.is_none()); + assert_eq!(count_total(&cfg).unwrap(), 1); + } +} diff --git a/src/openhuman/memory/tree/schemas.rs b/src/openhuman/memory/tree/schemas.rs index 6e5d59480e2..c558a85f668 100644 --- a/src/openhuman/memory/tree/schemas.rs +++ b/src/openhuman/memory/tree/schemas.rs @@ -23,6 +23,7 @@ pub fn all_controller_schemas() -> Vec { schemas("ingest"), schemas("list_chunks"), schemas("get_chunk"), + schemas("trigger_digest"), ] } @@ -40,6 +41,10 @@ pub fn all_registered_controllers() -> Vec { schema: schemas("get_chunk"), handler: handle_get_chunk, }, + RegisteredController { + schema: schemas("trigger_digest"), + handler: handle_trigger_digest, + }, ] } @@ -184,6 +189,44 @@ pub fn schemas(function: &str) -> ControllerSchema { required: false, }], }, + "trigger_digest" => ControllerSchema { + namespace: NAMESPACE, + function: "trigger_digest", + description: "Manually enqueue a daily-digest job for the global \ + tree. Idempotent — re-running for a day that already has a \ + digest is a no-op (the handler skips). When no date is \ + supplied, defaults to yesterday in UTC, matching the \ + scheduler's autonomous behavior.", + inputs: vec![FieldSchema { + name: "date_iso", + ty: TypeSchema::Option(Box::new(TypeSchema::String)), + comment: "UTC calendar date as `YYYY-MM-DD`. Optional; \ + defaults to yesterday when omitted.", + required: false, + }], + outputs: vec![ + FieldSchema { + name: "enqueued", + ty: TypeSchema::Bool, + comment: "True when a fresh job row was inserted; false \ + when an active job for the same date suppressed it.", + required: true, + }, + FieldSchema { + name: "job_id", + ty: TypeSchema::Option(Box::new(TypeSchema::String)), + comment: "ID of the newly enqueued job row, when enqueued.", + required: false, + }, + FieldSchema { + name: "date_iso", + ty: TypeSchema::String, + comment: "The date the digest will cover, echoed back \ + as `YYYY-MM-DD`.", + required: true, + }, + ], + }, _ => ControllerSchema { namespace: NAMESPACE, function: "unknown", @@ -228,6 +271,14 @@ fn handle_get_chunk(params: Map) -> ControllerFuture { }) } +fn handle_trigger_digest(params: Map) -> ControllerFuture { + Box::pin(async move { + let config = config_rpc::load_config_with_timeout().await?; + let req = parse_value::(Value::Object(params))?; + to_json(tree_rpc::trigger_digest_rpc(&config, req).await?) + }) +} + fn parse_value(v: Value) -> Result { serde_json::from_value(v).map_err(|e| format!("invalid params: {e}")) } diff --git a/src/openhuman/memory/tree/score/extract/llm.rs b/src/openhuman/memory/tree/score/extract/llm.rs index 6ff895618ba..0ee2469c9b2 100644 --- a/src/openhuman/memory/tree/score/extract/llm.rs +++ b/src/openhuman/memory/tree/score/extract/llm.rs @@ -36,7 +36,7 @@ use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; -use super::types::{EntityKind, ExtractedEntities, ExtractedEntity}; +use super::types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; use super::EntityExtractor; // ── Configuration ──────────────────────────────────────────────────────── @@ -59,6 +59,15 @@ pub struct LlmExtractorConfig { /// If true, drop entities whose declared kind isn't in `allowed_kinds` /// instead of falling back to [`EntityKind::Misc`]. pub strict_kinds: bool, + /// If true, the system prompt asks the model to also emit a + /// `topics` array (free-form theme labels), and the response parser + /// populates [`ExtractedEntities::topics`]. Default `false` — the + /// extractor's primary job is named-entity extraction; topics are + /// an opt-in side-channel for callers that need a thematic + /// summary in the same call (e.g. running over a sealed summary's + /// content). Adds prompt tokens and gives the model one more + /// schema field to keep track of, so leave off unless needed. + pub emit_topics: bool, } impl Default for LlmExtractorConfig { @@ -73,8 +82,13 @@ impl Default for LlmExtractorConfig { EntityKind::Location, EntityKind::Event, EntityKind::Product, + EntityKind::Datetime, + EntityKind::Technology, + EntityKind::Artifact, + EntityKind::Quantity, ], strict_kinds: false, + emit_topics: false, } } } @@ -103,7 +117,7 @@ impl LlmEntityExtractor { messages: vec![ OllamaMessage { role: "system".to_string(), - content: SYSTEM_PROMPT.to_string(), + content: build_system_prompt(self.cfg.emit_topics), }, OllamaMessage { role: "user".to_string(), @@ -204,28 +218,87 @@ impl LlmEntityExtractor { // ── Prompt ─────────────────────────────────────────────────────────────── -const SYSTEM_PROMPT: &str = "\ -You are a named-entity extractor and importance rater. Return JSON only — \ +/// Build the system prompt for the extractor. When `emit_topics` is true +/// the schema, required-fields list, and example outputs include a +/// `topics` array (free-form theme labels). When false the prompt +/// matches the pre-flag behaviour exactly — no mention of topics +/// anywhere — so the small model isn't asked to produce a field the +/// caller doesn't want. +fn build_system_prompt(emit_topics: bool) -> String { + let topics_schema_line = if emit_topics { + " \"topics\": [\"\"],\n" + } else { + "" + }; + let topics_required = if emit_topics { "topics, " } else { "" }; + let fields_count = if emit_topics { "four" } else { "three" }; + let topics_guide = if emit_topics { + "Topics are short free-form theme labels for what the text is ABOUT \ + (e.g. \"rate limiting\", \"memory tree\", \"auth flow\"). They are \ + distinct from entities — entities are specific named things mentioned \ + in the text; topics are the abstract themes those things relate to.\n" + } else { + "" + }; + let example1_topics = if emit_topics { + ",\"topics\":[\"shipping\",\"auth\"]" + } else { + "" + }; + let example2_topics = if emit_topics { + ",\"topics\":[\"product launch\",\"revenue\"]" + } else { + "" + }; + + format!( + "You are a named-entity extractor and importance rater. Return JSON only — \ no prose, no markdown, no commentary. Do not summarize. Extract every named \ entity mention you find, including duplicates, and rate the chunk's overall \ importance as a float in [0.0, 1.0]. Schema: -{ +{{ \"entities\": [ - { \"kind\": \"person|organization|location|event|product\", - \"text\": \"\" } + {{ \"kind\": \"person|organization|location|event|product|datetime|technology|artifact|quantity\", + \"text\": \"\" }} ], - \"importance\": 0.0, +{topics_schema_line} \"importance\": 0.0, \"importance_reason\": \"\" -} +}} + +Kinds guide: + person named human (\"Alice\", \"Steven Enamakel\") + organization company / team / project (\"Anthropic\", \"TinyHumans\") + location place (\"SF office\", \"London\") + event scheduled occurrence (\"Q2 launch\", \"design review\") + product commercial offering (\"Claude Code\", \"OpenHuman\") + datetime temporal expression (\"Friday\", \"Q2 2026\", \"EOD tomorrow\") + technology tool / framework / language / service (\"Rust\", \"OAuth\", \"Slack API\") + artifact code / ticket / doc reference (\"PR #934\", \"src/foo.rs\", \"OH-42\") + quantity amount / metric / money (\"$5K\", \"20/min\", \"10k tokens\") + +{topics_guide} +If a mention doesn't clearly fit a kind above, omit it rather than guessing. +Always emit ALL {fields_count} top-level fields (entities, {topics_required}importance, importance_reason), +even when entities is empty. + +Examples: + +Input: alice and bob shipped the auth migration friday. PR #42 ships OAuth refactor in src/auth/. +Output: {{\"entities\":[{{\"kind\":\"person\",\"text\":\"alice\"}},{{\"kind\":\"person\",\"text\":\"bob\"}},{{\"kind\":\"event\",\"text\":\"auth migration\"}},{{\"kind\":\"datetime\",\"text\":\"friday\"}},{{\"kind\":\"artifact\",\"text\":\"PR #42\"}},{{\"kind\":\"technology\",\"text\":\"OAuth\"}},{{\"kind\":\"artifact\",\"text\":\"src/auth/\"}}]{example1_topics},\"importance\":0.9,\"importance_reason\":\"explicit shipping commitment\"}} + +Input: Anthropic shipped Claude Code in SF — $20M ARR target by Q2. +Output: {{\"entities\":[{{\"kind\":\"organization\",\"text\":\"Anthropic\"}},{{\"kind\":\"product\",\"text\":\"Claude Code\"}},{{\"kind\":\"location\",\"text\":\"SF\"}},{{\"kind\":\"quantity\",\"text\":\"$20M ARR\"}},{{\"kind\":\"datetime\",\"text\":\"Q2\"}}]{example2_topics},\"importance\":0.85,\"importance_reason\":\"factual content with key business metric\"}} Importance guide: 0.9+ actionable decisions, key information, explicit commitments 0.6+ substantive discussion, factual content, named entities 0.3+ ambient context, low-density prose <0.3 reactions, acknowledgments, bots, trivial exchanges -"; +" + ) +} // ── Wire types (Ollama API) ────────────────────────────────────────────── @@ -265,6 +338,11 @@ struct OllamaResponseMessage { struct LlmExtractionOutput { #[serde(default)] entities: Vec, + /// Free-form theme labels — populated only when the extractor is + /// configured with `emit_topics = true`. Always tolerant of absence + /// so models that ignore the field don't fail parsing. + #[serde(default)] + topics: Vec, #[serde(default)] importance: Option, #[serde(default)] @@ -355,9 +433,26 @@ impl LlmExtractionOutput { let llm_importance = self.importance.map(|v| v.clamp(0.0, 1.0)); + // Topics: only populated when the caller enabled `emit_topics` + // (the prompt asked for them). Otherwise this is empty by + // default — the model didn't know to emit topics, so any value + // here would be hallucination. + let topics = self + .topics + .into_iter() + .filter_map(|raw| { + let label = raw.trim().to_string(); + if label.is_empty() { + None + } else { + Some(ExtractedTopic { label, score: 0.85 }) + } + }) + .collect(); + ExtractedEntities { entities, - topics: Vec::new(), + topics, llm_importance, llm_importance_reason: self.importance_reason, } @@ -373,6 +468,14 @@ fn parse_kind(s: &str) -> Option { "location" | "place" | "loc" => Some(EntityKind::Location), "event" => Some(EntityKind::Event), "product" => Some(EntityKind::Product), + "datetime" | "date" | "time" | "timestamp" => Some(EntityKind::Datetime), + "technology" | "tech" | "tool" | "framework" | "library" | "language" | "service" => { + Some(EntityKind::Technology) + } + "artifact" | "reference" | "ref" | "pr" | "ticket" | "file" | "commit" => { + Some(EntityKind::Artifact) + } + "quantity" | "amount" | "metric" | "number" | "money" => Some(EntityKind::Quantity), "misc" | "miscellaneous" | "other" => Some(EntityKind::Misc), _ => None, } diff --git a/src/openhuman/memory/tree/score/extract/llm_tests.rs b/src/openhuman/memory/tree/score/extract/llm_tests.rs index a1f61c5b2fd..0ee9bf36441 100644 --- a/src/openhuman/memory/tree/score/extract/llm_tests.rs +++ b/src/openhuman/memory/tree/score/extract/llm_tests.rs @@ -1,5 +1,39 @@ use super::*; +#[test] +fn build_system_prompt_default_omits_topics() { + let p = build_system_prompt(false); + assert!(!p.contains("\"topics\"")); + assert!(!p.contains("Topics are")); + assert!(p.contains("ALL three top-level fields")); + assert!(p.contains("entities, importance")); +} + +#[test] +fn build_system_prompt_with_flag_includes_topics() { + let p = build_system_prompt(true); + assert!(p.contains("\"topics\"")); + assert!(p.contains("Topics are short free-form theme labels")); + assert!(p.contains("ALL four top-level fields")); + assert!(p.contains("entities, topics, importance")); +} + +#[test] +fn extraction_output_parses_topics_when_present() { + let json = r#"{"entities":[],"topics":["rate limiting","memory tree"],"importance":0.6,"importance_reason":"r"}"#; + let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap(); + assert_eq!(parsed.topics, vec!["rate limiting", "memory tree"]); +} + +#[test] +fn extraction_output_tolerates_missing_topics() { + // Default extractor (emit_topics=false) — model won't emit topics + // and parsing must still succeed. + let json = r#"{"entities":[],"importance":0.6,"importance_reason":"r"}"#; + let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap(); + assert!(parsed.topics.is_empty()); +} + #[test] fn parse_kind_normalisation() { assert_eq!(parse_kind("Person"), Some(EntityKind::Person)); @@ -8,6 +42,42 @@ fn parse_kind_normalisation() { assert!(parse_kind("Spaceship").is_none()); } +#[test] +fn parse_kind_accepts_new_semantic_kinds_and_synonyms() { + // Datetime + for s in ["datetime", "date", "time", "timestamp", " DateTime "] { + assert_eq!(parse_kind(s), Some(EntityKind::Datetime), "input={s:?}"); + } + // Technology + for s in [ + "technology", + "tech", + "tool", + "framework", + "library", + "language", + "service", + ] { + assert_eq!(parse_kind(s), Some(EntityKind::Technology), "input={s:?}"); + } + // Artifact + for s in [ + "artifact", + "reference", + "ref", + "pr", + "ticket", + "file", + "commit", + ] { + assert_eq!(parse_kind(s), Some(EntityKind::Artifact), "input={s:?}"); + } + // Quantity + for s in ["quantity", "amount", "metric", "number", "money"] { + assert_eq!(parse_kind(s), Some(EntityKind::Quantity), "input={s:?}"); + } +} + #[test] fn find_char_span_handles_unicode() { let text = "中 Alice met Bob"; @@ -71,6 +141,7 @@ fn into_extracted_entities_gives_distinct_spans_to_duplicate_mentions() { text: "Alice".into(), }, ], + topics: vec![], importance: None, importance_reason: None, }; @@ -100,6 +171,7 @@ fn into_extracted_entities_drops_extra_duplicate_when_source_only_has_one() { text: "Alice".into(), }, ], + topics: vec![], importance: None, importance_reason: None, }; @@ -138,6 +210,7 @@ fn into_extracted_entities_drops_hallucinations() { text: "ImaginaryPerson".into(), }, ], + topics: vec![], importance: Some(0.7), importance_reason: Some("substantive".into()), }; @@ -154,6 +227,7 @@ fn into_extracted_entities_drops_hallucinations() { fn into_extracted_entities_clamps_importance() { let out = LlmExtractionOutput { entities: vec![], + topics: vec![], importance: Some(1.5), importance_reason: None, }; @@ -169,6 +243,7 @@ fn into_extracted_entities_strict_drops_unknown_kinds() { kind: "spaceship".into(), text: "Enterprise".into(), }], + topics: vec![], importance: None, importance_reason: None, }; @@ -187,6 +262,7 @@ fn into_extracted_entities_lenient_falls_back_to_misc() { kind: "spaceship".into(), text: "Enterprise".into(), }], + topics: vec![], importance: None, importance_reason: None, }; @@ -204,6 +280,7 @@ fn into_extracted_entities_disallowed_known_kind_falls_back_to_misc() { kind: "person".into(), text: "Alice".into(), }], + topics: vec![], importance: None, importance_reason: None, }; diff --git a/src/openhuman/memory/tree/score/extract/mod.rs b/src/openhuman/memory/tree/score/extract/mod.rs index f79a97b197b..e21beda8471 100644 --- a/src/openhuman/memory/tree/score/extract/mod.rs +++ b/src/openhuman/memory/tree/score/extract/mod.rs @@ -10,6 +10,80 @@ pub mod llm; pub mod regex; pub mod types; +use std::sync::Arc; + +use crate::openhuman::config::Config; + pub use extractor::{CompositeExtractor, EntityExtractor, RegexEntityExtractor}; pub use llm::{LlmEntityExtractor, LlmExtractorConfig}; pub use types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic}; + +/// Build the extractor used by seal handlers to label new summary nodes. +/// +/// Composition: +/// - regex extractor — always on, mechanical, near-zero cost +/// - LLM extractor with `emit_topics: true` — added when +/// `memory_tree.llm_extractor_endpoint` and `..._model` are both set +/// +/// Differs from [`super::ScoringConfig::from_config`] (the chunk-admission +/// builder) in two ways: returns *just* an extractor (no thresholds / +/// weights / drop logic — none of which apply at seal time), and flips +/// `emit_topics` on so summaries surface thematic labels alongside +/// entities. Leaf-side scoring is unchanged. +pub fn build_summary_extractor(config: &Config) -> Arc { + let endpoint = config + .memory_tree + .llm_extractor_endpoint + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()); + let model = config + .memory_tree + .llm_extractor_model + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()); + + let (Some(endpoint), Some(model)) = (endpoint, model) else { + log::debug!( + "[memory_tree::extract] summary extractor: LLM not configured — using regex-only" + ); + return Arc::new(CompositeExtractor::regex_only()); + }; + + let timeout_ms = config + .memory_tree + .llm_extractor_timeout_ms + .unwrap_or(15_000); + + let cfg = LlmExtractorConfig { + endpoint: endpoint.to_string(), + model: model.to_string(), + timeout: std::time::Duration::from_millis(timeout_ms), + emit_topics: true, + ..LlmExtractorConfig::default() + }; + + match LlmEntityExtractor::new(cfg) { + Ok(llm) => { + log::info!( + "[memory_tree::extract] summary extractor: regex + LLM endpoint={} model={} \ + timeout_ms={} emit_topics=true", + endpoint, + model, + timeout_ms + ); + Arc::new(CompositeExtractor::new(vec![ + Box::new(RegexEntityExtractor), + Box::new(llm), + ])) + } + Err(err) => { + log::warn!( + "[memory_tree::extract] summary extractor: LlmEntityExtractor construction \ + failed: {err:#} — falling back to regex-only" + ); + Arc::new(CompositeExtractor::regex_only()) + } + } +} diff --git a/src/openhuman/memory/tree/score/extract/types.rs b/src/openhuman/memory/tree/score/extract/types.rs index 8e91c843718..89b4fd6726f 100644 --- a/src/openhuman/memory/tree/score/extract/types.rs +++ b/src/openhuman/memory/tree/score/extract/types.rs @@ -26,12 +26,22 @@ pub enum EntityKind { Url, Handle, Hashtag, - // Semantic (reserved — not emitted in Phase 2) + // Semantic — emitted by the LLM extractor. Person, Organization, Location, Event, Product, + /// Temporal expressions: "Friday", "Q2 2026", "EOD tomorrow", "next sprint". + Datetime, + /// Tools / frameworks / programming languages / services: + /// "Rust", "OAuth", "Slack API", "nomic-embed". + Technology, + /// Code / ticket / doc references that point at something addressable: + /// "PR #934", "src/openhuman/...", "OH-42", "ab7da2e2". + Artifact, + /// Amounts / metrics / money: "$5K", "20/min", "10k tokens", "52 chunks". + Quantity, Misc, // Thematic — scorer-surfaced topics (hashtag-like short phrases or // LLM-extracted themes). Promoted into the canonical entity stream @@ -54,6 +64,10 @@ impl EntityKind { Self::Location => "location", Self::Event => "event", Self::Product => "product", + Self::Datetime => "datetime", + Self::Technology => "technology", + Self::Artifact => "artifact", + Self::Quantity => "quantity", Self::Misc => "misc", Self::Topic => "topic", } @@ -70,6 +84,10 @@ impl EntityKind { "location" => Ok(Self::Location), "event" => Ok(Self::Event), "product" => Ok(Self::Product), + "datetime" => Ok(Self::Datetime), + "technology" => Ok(Self::Technology), + "artifact" => Ok(Self::Artifact), + "quantity" => Ok(Self::Quantity), "misc" => Ok(Self::Misc), "topic" => Ok(Self::Topic), other => Err(format!("unknown entity kind: {other}")), @@ -207,7 +225,12 @@ mod tests { EntityKind::Location, EntityKind::Event, EntityKind::Product, + EntityKind::Datetime, + EntityKind::Technology, + EntityKind::Artifact, + EntityKind::Quantity, EntityKind::Misc, + EntityKind::Topic, ] { assert_eq!(EntityKind::parse(k.as_str()).unwrap(), k); } diff --git a/src/openhuman/memory/tree/score/mod.rs b/src/openhuman/memory/tree/score/mod.rs index 15f17be767c..7ae5ade51a7 100644 --- a/src/openhuman/memory/tree/score/mod.rs +++ b/src/openhuman/memory/tree/score/mod.rs @@ -336,6 +336,22 @@ pub async fn score_chunks(chunks: &[Chunk], cfg: &ScoringConfig) -> Result Result> { + let fast_cfg = ScoringConfig { + extractor: cfg.extractor.clone(), + weights: cfg.weights.clone(), + drop_threshold: cfg.drop_threshold, + llm_extractor: None, + definite_keep_threshold: cfg.definite_keep_threshold, + definite_drop_threshold: cfg.definite_drop_threshold, + }; + score_chunks(chunks, &fast_cfg).await +} + // ── Persistence helpers used by the ingest orchestrator ───────────────── /// Persist the score row + entity-index rows for one kept chunk. diff --git a/src/openhuman/memory/tree/score/store.rs b/src/openhuman/memory/tree/score/store.rs index ecb314b8da2..e8670e57bcb 100644 --- a/src/openhuman/memory/tree/score/store.rs +++ b/src/openhuman/memory/tree/score/store.rs @@ -371,6 +371,21 @@ pub fn lookup_entity( }) } +pub fn list_entity_ids_for_node(config: &Config, node_id: &str) -> Result> { + with_connection(config, |conn| { + let mut stmt = conn.prepare( + "SELECT DISTINCT entity_id + FROM mem_tree_entity_index + WHERE node_id = ?1 + ORDER BY score DESC, timestamp_ms DESC, entity_id ASC", + )?; + let rows = stmt + .query_map(params![node_id], |row| row.get::<_, String>(0))? + .collect::>>()?; + Ok(rows) + }) +} + /// Count rows in the entity index (for tests / diagnostics). pub fn count_entity_index(config: &Config) -> Result { with_connection(config, |conn| { diff --git a/src/openhuman/memory/tree/source_tree/bucket_seal.rs b/src/openhuman/memory/tree/source_tree/bucket_seal.rs index 0f9f2b2c0b4..44445a31adb 100644 --- a/src/openhuman/memory/tree/source_tree/bucket_seal.rs +++ b/src/openhuman/memory/tree/source_tree/bucket_seal.rs @@ -26,12 +26,17 @@ //! summariser does no real I/O; when a networked summariser lands, wrap //! DB calls in `tokio::task::spawn_blocking` to keep the runtime healthy. +use std::collections::BTreeSet; +use std::sync::Arc; + use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; use rusqlite::Transaction; use crate::openhuman::config::Config; use crate::openhuman::memory::tree::score::embed::build_embedder_from_config; +use crate::openhuman::memory::tree::score::extract::EntityExtractor; +use crate::openhuman::memory::tree::score::resolver::canonicalise; use crate::openhuman::memory::tree::source_tree::registry::new_summary_id; use crate::openhuman::memory::tree::source_tree::store; use crate::openhuman::memory::tree::source_tree::summariser::{ @@ -47,6 +52,88 @@ use crate::openhuman::memory::tree::store::with_connection; /// realistic source. const MAX_CASCADE_DEPTH: u32 = 32; +/// How a sealed summary node's `entities` and `topics` fields get populated. +/// +/// Each tree kind has different correct semantics: +/// - **Source** trees use [`LabelStrategy::ExtractFromContent`] so the +/// summariser's freshly-synthesised text gets its own pass through an +/// extractor. Captures emergent themes that no individual leaf expressed. +/// - **Global** trees use [`LabelStrategy::UnionFromChildren`] — their +/// inputs are already-labeled source-tree summaries; union preserves +/// labels for time-based retrieval ("days that mentioned Alice") +/// without an LLM call. +/// - **Topic** trees use [`LabelStrategy::Empty`] — their scope already +/// pins the dominant theme; inheriting auxiliary entities would +/// cross-pollinate unrelated topic trees and noise the entity index. +#[derive(Clone)] +pub enum LabelStrategy { + /// Run the extractor on the new summary's content; canonicalise the + /// result into `entities` (canonical_ids) and `topics` (labels). + ExtractFromContent(Arc), + /// Dedup-merge each input's `entities` and `topics` into the parent. + UnionFromChildren, + /// Leave both fields empty regardless of inputs. + Empty, +} + +impl std::fmt::Debug for LabelStrategy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::ExtractFromContent(ex) => write!(f, "ExtractFromContent({})", ex.name()), + Self::UnionFromChildren => f.write_str("UnionFromChildren"), + Self::Empty => f.write_str("Empty"), + } + } +} + +/// Resolve `entities` and `topics` for a freshly-summarised node according +/// to the chosen strategy. Errors propagate from the extractor (when used). +async fn resolve_labels( + strategy: &LabelStrategy, + inputs: &[SummaryInput], + summary_content: &str, +) -> Result<(Vec, Vec)> { + match strategy { + LabelStrategy::ExtractFromContent(extractor) => { + let extracted = extractor + .extract(summary_content) + .await + .context("seal-time extractor failed")?; + let canonical = canonicalise(&extracted); + let mut entities: Vec = canonical + .into_iter() + .map(|c| c.canonical_id) + .collect::>() + .into_iter() + .collect(); + entities.sort(); + let mut topics: Vec = extracted + .topics + .into_iter() + .map(|t| t.label) + .collect::>() + .into_iter() + .collect(); + topics.sort(); + Ok((entities, topics)) + } + LabelStrategy::UnionFromChildren => { + let mut entities: BTreeSet = BTreeSet::new(); + let mut topics: BTreeSet = BTreeSet::new(); + for inp in inputs { + for e in &inp.entities { + entities.insert(e.clone()); + } + for t in &inp.topics { + topics.insert(t.clone()); + } + } + Ok((entities.into_iter().collect(), topics.into_iter().collect())) + } + LabelStrategy::Empty => Ok((Vec::new(), Vec::new())), + } +} + /// A single leaf being appended to an L0 buffer. #[derive(Clone, Debug)] pub struct LeafRef { @@ -61,17 +148,22 @@ pub struct LeafRef { /// Append a leaf to the source tree for `tree`, sealing buffers as they /// fill. Returns the ids of any summaries that sealed during this call. +/// +/// `strategy` controls how each sealed summary's `entities` and `topics` +/// are populated — see [`LabelStrategy`]. pub async fn append_leaf( config: &Config, tree: &Tree, leaf: &LeafRef, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result> { log::debug!( - "[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={}", + "[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={} strategy={:?}", tree.id, leaf.chunk_id, - leaf.token_count + leaf.token_count, + strategy ); // 1. Push leaf into L0 buffer (transactional). @@ -85,7 +177,24 @@ pub async fn append_leaf( )?; // 2. Cascade seals upward until a level stays under budget. - cascade_seals(config, tree, summariser).await + cascade_seals(config, tree, summariser, strategy).await +} + +/// Queue-oriented variant of [`append_leaf`]. +/// +/// This only appends the leaf to the L0 buffer and returns whether the +/// caller should enqueue a follow-up seal job for level 0. +pub fn append_leaf_deferred(config: &Config, tree: &Tree, leaf: &LeafRef) -> Result { + append_to_buffer( + config, + &tree.id, + 0, + &leaf.chunk_id, + leaf.token_count as i64, + leaf.timestamp, + )?; + let buf = store::get_buffer(config, &tree.id, 0)?; + Ok(should_seal(&buf)) } /// Transactionally append a single item to `(tree_id, level)`'s buffer. @@ -127,20 +236,25 @@ async fn cascade_seals( config: &Config, tree: &Tree, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result> { - cascade_all_from(config, tree, 0, summariser, None).await + cascade_all_from(config, tree, 0, summariser, None, strategy).await } /// Seal buffers starting at `start_level` and cascade upward. When /// `force_now` is `Some`, the buffer at `start_level` is sealed regardless /// of token budget (used by time-based flush). Upper levels are sealed /// only when they cross the budget. +/// +/// `strategy` is forwarded to every sealed level — same semantics as +/// [`append_leaf`]. pub async fn cascade_all_from( config: &Config, tree: &Tree, start_level: u32, summariser: &dyn Summariser, force_now: Option>, + strategy: &LabelStrategy, ) -> Result> { let mut sealed_ids: Vec = Vec::new(); let mut level: u32 = start_level; @@ -169,7 +283,9 @@ pub async fn cascade_all_from( break; } - let summary_id = seal_one_level(config, tree, &buf, summariser).await?; + // Sync cascade — drives the level walk itself; doesn't need the + // queue follow-ups (we'll hit `seal_one_level` again next iter). + let summary_id = seal_one_level(config, tree, &buf, summariser, strategy, false).await?; sealed_ids.push(summary_id); level += 1; } @@ -185,7 +301,7 @@ pub async fn cascade_all_from( /// summarisers that emit at the full token budget (e.g. the inert /// fallback) collapse the cascade into a 1:1:1 chain instead of a real /// tree. -fn should_seal(buf: &Buffer) -> bool { +pub(crate) fn should_seal(buf: &Buffer) -> bool { if buf.is_empty() { return false; } @@ -198,11 +314,31 @@ fn should_seal(buf: &Buffer) -> bool { /// Seal `buf` at `level` into one summary at `level + 1`. Returns the new /// summary id. -async fn seal_one_level( +/// +/// `strategy` decides how `entities` and `topics` get populated on the new +/// summary node — see [`LabelStrategy`]. +/// +/// When `enqueue_follow_ups` is `true`, the function additionally inserts +/// follow-up job rows **inside the same transaction** that commits the +/// seal: +/// - `seal { tree_id, level: parent_level }` if the parent buffer's gate +/// is now met (parent-cascade enqueue) +/// - `topic_route { NodeRef::Summary { summary_id } }` for source trees +/// (so summary-level entities feed the topic-tree spawn pipeline) +/// +/// Atomic enqueue eliminates the crash window where a seal commits but +/// the post-commit follow-up enqueues are silently lost on a worker +/// crash. The async-pipeline handler (`handle_seal`) passes `true`. The +/// synchronous in-process cascade caller ([`cascade_all_from`]) passes +/// `false` because it drives the cascade itself and topic_route isn't +/// part of the test/flush sync path. +pub(crate) async fn seal_one_level( config: &Config, tree: &Tree, buf: &Buffer, summariser: &dyn Summariser, + strategy: &LabelStrategy, + enqueue_follow_ups: bool, ) -> Result { let level = buf.level; let target_level = level + 1; @@ -237,7 +373,7 @@ async fn seal_one_level( // Run summariser — async, OUTSIDE any DB transaction. let ctx = SummaryContext { tree_id: &tree.id, - tree_kind: TreeKind::Source, + tree_kind: tree.kind, target_level, token_budget: TOKEN_BUDGET, }; @@ -246,6 +382,12 @@ async fn seal_one_level( .await .context("summariser failed during seal")?; + // Resolve labels (entities/topics) for the new summary node according + // to the chosen strategy. Done before the write tx so an extractor + // failure aborts the seal cleanly — same shape as the embedder guard + // below. + let (node_entities, node_topics) = resolve_labels(strategy, &inputs, &output.content).await?; + // Phase 4 (#710): embed the summary BEFORE opening the write tx so an // embedder failure aborts the seal cleanly — nothing is persisted, // the buffer stays intact, and a retry re-embeds from scratch. The @@ -291,14 +433,18 @@ async fn seal_one_level( let node = SummaryNode { id: summary_id.clone(), tree_id: tree.id.clone(), - tree_kind: TreeKind::Source, + // `seal_one_level` runs for source AND topic trees (handle_seal, + // cascade_all_from, flush). Hardcoding Source here would write + // topic-tree summaries with tree_kind='source' in + // mem_tree_summaries, breaking any query filtering on tree_kind. + tree_kind: tree.kind, level: target_level, parent_id: None, child_ids: buf.item_ids.clone(), content: output.content, token_count: output.token_count, - entities: output.entities, - topics: output.topics, + entities: node_entities, + topics: node_topics, time_range_start, time_range_end, score, @@ -308,12 +454,15 @@ async fn seal_one_level( }; // Single write transaction: insert summary, clear this buffer, append - // summary id to parent buffer, bump tree max_level/root if needed. + // summary id to parent buffer, bump tree max_level/root if needed, + // and (when `enqueue_follow_ups`) atomically enqueue parent-seal + + // topic_route follow-ups so they can never desync from the commit. // Re-read `max_level` from inside the tx so cascading seals within // one call see the updated value from earlier levels. let summary_id_for_closure = summary_id.clone(); let target_level_for_closure = target_level; let tree_id = tree.id.clone(); + let tree_kind = tree.kind; with_connection(config, move |conn| { let tx = conn.unchecked_transaction()?; @@ -375,6 +524,43 @@ async fn seal_one_level( }; store::upsert_buffer_tx(&tx, &parent)?; + // Atomic follow-up enqueues. Done INSIDE this tx — if the commit + // rolls back, the queue rows go with it; if it succeeds, the + // rows are durably visible to the worker pool. Eliminates the + // crash window where the seal commits but post-commit enqueues + // are lost. + if enqueue_follow_ups { + // Parent-cascade: if the new summary made the parent buffer + // cross its gate, enqueue the next level's seal. Dedupe key + // `seal:{tree_id}:{parent_level}` prevents duplicates if a + // parallel path already queued it. + if should_seal(&parent) { + use crate::openhuman::memory::tree::jobs::store::enqueue_tx as enqueue_job_tx; + use crate::openhuman::memory::tree::jobs::types::{NewJob, SealPayload}; + let parent_seal = SealPayload { + tree_id: tree_id.clone(), + level: target_level_for_closure, + force_now_ms: None, + }; + enqueue_job_tx(&tx, &NewJob::seal(&parent_seal)?)?; + } + // Source-tree summary routing: feed the new summary's + // entities back into the topic-tree spawn pipeline. Topic + // and global trees are sinks — no fan-out from their seals. + if matches!(tree_kind, TreeKind::Source) { + use crate::openhuman::memory::tree::jobs::store::enqueue_tx as enqueue_job_tx; + use crate::openhuman::memory::tree::jobs::types::{ + NewJob, NodeRef, TopicRoutePayload, + }; + let route = TopicRoutePayload { + node: NodeRef::Summary { + summary_id: summary_id_for_closure.clone(), + }, + }; + enqueue_job_tx(&tx, &NewJob::topic_route(&route)?)?; + } + } + // Update tree root / max_level if we just climbed. if target_level_for_closure > current_max { store::update_tree_after_seal_tx( @@ -447,7 +633,7 @@ fn hydrate_inputs(config: &Config, level: u32, item_ids: &[String]) -> Result Result> { - use crate::openhuman::memory::tree::score::store::get_score; + use crate::openhuman::memory::tree::score::store::{get_score, list_entity_ids_for_node}; use crate::openhuman::memory::tree::store::get_chunk; let mut out: Vec = Vec::with_capacity(chunk_ids.len()); @@ -461,17 +647,19 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result (row.total, Vec::new(), chunk.metadata.tags.clone()), - None => (0.0, Vec::new(), chunk.metadata.tags.clone()), - }; + let score_value = get_score(config, id)?.map(|row| row.total).unwrap_or(0.0); + // Pull canonical entity ids from the inverted index — that's the + // authoritative source for "what entities are attached to this + // chunk." Topics live on the chunk's metadata tags. + // [`LabelStrategy::UnionFromChildren`] reads these fields off + // each `SummaryInput` to roll labels up the tree. + let entities = list_entity_ids_for_node(config, id).unwrap_or_default(); out.push(SummaryInput { id: chunk.id.clone(), content: chunk.content.clone(), token_count: chunk.token_count, entities, - topics, + topics: chunk.metadata.tags.clone(), time_range_start: chunk.metadata.time_range.0, time_range_end: chunk.metadata.time_range.1, score: score_value, diff --git a/src/openhuman/memory/tree/source_tree/bucket_seal_tests.rs b/src/openhuman/memory/tree/source_tree/bucket_seal_tests.rs index a39524ba431..d3b4fdad66b 100644 --- a/src/openhuman/memory/tree/source_tree/bucket_seal_tests.rs +++ b/src/openhuman/memory/tree/source_tree/bucket_seal_tests.rs @@ -36,7 +36,9 @@ async fn append_below_budget_does_not_seal() { // Chunks don't exist in DB — we're only exercising the buffer // accounting, which doesn't require leaf rows until a seal fires. let leaf = mk_leaf("leaf-1", 100, 1_700_000_000_000); - let sealed = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert!(sealed.is_empty()); let buf = store::get_buffer(&cfg, &tree.id, 0).unwrap(); @@ -97,10 +99,14 @@ async fn crossing_budget_triggers_seal() { score: 0.5, }; - let first = append_leaf(&cfg, &tree, &leaf1, &summariser).await.unwrap(); + let first = append_leaf(&cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert!(first.is_empty(), "first append below budget — no seal"); - let second = append_leaf(&cfg, &tree, &leaf2, &summariser).await.unwrap(); + let second = append_leaf(&cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert_eq!(second.len(), 1, "second append crosses budget — one seal"); let summary_id = &second[0]; @@ -186,7 +192,9 @@ async fn fanout_at_l1_triggers_l2_seal() { topics: vec![], score: 0.5, }; - let sealed = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); all_sealed.extend(sealed); } @@ -264,7 +272,9 @@ async fn upper_level_does_not_seal_below_fanout() { topics: vec![], score: 0.5, }; - let _ = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); + let _ = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); } let t = store::get_tree(&cfg, &tree.id).unwrap().unwrap(); @@ -281,3 +291,279 @@ async fn upper_level_does_not_seal_below_fanout() { stop_before as u64 ); } + +// ── LabelStrategy tests (#TBD) ──────────────────────────────────────────── +// +// These exercise the three labeling modes seal_one_level supports. We use +// a short token budget so the seal fires on a single leaf — keeps the +// arithmetic of "what entities/topics end up on the parent" obvious. + +/// Helper: persist a substantive chunk and return a `LeafRef` referencing +/// it, with caller-supplied entity/topic labels (used by Union/Empty tests). +/// +/// To match production, entity labels are written into `mem_tree_entity_index` +/// (where seal-time hydration reads them from) and topic labels are stored +/// on `chunk.metadata.tags` (the production source of leaf-level topics). +fn seed_leaf( + cfg: &Config, + seq: u32, + content: &str, + entities: Vec, + topics: Vec, +) -> LeafRef { + use crate::openhuman::memory::tree::score::extract::EntityKind; + use crate::openhuman::memory::tree::score::resolver::CanonicalEntity; + use crate::openhuman::memory::tree::score::store::index_entity; + use crate::openhuman::memory::tree::store::upsert_chunks; + use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef}; + use chrono::TimeZone; + let ts = Utc + .timestamp_millis_opt(1_700_000_000_000 + seq as i64) + .unwrap(); + let chunk = Chunk { + id: chunk_id(SourceKind::Chat, "slack:#eng", seq, content), + content: content.to_string(), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: "slack:#eng".into(), + owner: "alice".into(), + timestamp: ts, + time_range: (ts, ts), + tags: topics.clone(), + source_ref: Some(SourceRef::new(format!("slack://x{seq}"))), + }, + // Bust TOKEN_BUDGET in one leaf so the seal fires immediately. + token_count: 10_000, + seq_in_source: seq, + created_at: ts, + }; + upsert_chunks(cfg, &[chunk.clone()]).unwrap(); + // Mirror production indexing: entities go into mem_tree_entity_index + // so the seal hydrator can pull them via list_entity_ids_for_node. + for entity_id in &entities { + let kind = entity_id + .split_once(':') + .map_or(EntityKind::Misc, |(k, _)| { + EntityKind::parse(k).unwrap_or(EntityKind::Misc) + }); + let surface = entity_id + .split_once(':') + .map_or(entity_id.as_str(), |(_, v)| v); + let e = CanonicalEntity { + canonical_id: entity_id.clone(), + kind, + surface: surface.to_string(), + span_start: 0, + span_end: surface.len() as u32, + score: 1.0, + }; + index_entity(cfg, &e, &chunk.id, "leaf", ts.timestamp_millis(), None).unwrap(); + } + LeafRef { + chunk_id: chunk.id.clone(), + token_count: chunk.token_count, + timestamp: ts, + content: chunk.content.clone(), + entities, + topics, + score: 0.5, + } +} + +#[tokio::test] +async fn seal_with_extract_strategy_populates_entities_and_topics() { + use crate::openhuman::memory::tree::score::extract::{CompositeExtractor, EntityExtractor}; + use std::sync::Arc; + + let (_tmp, cfg) = test_config(); + let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); + let summariser = InertSummariser::new(); + + // Content the regex extractor can find: an email and a hashtag. The + // inert summariser concatenates leaf content into the L1 summary, so + // these tokens survive into the summary text and the extractor finds + // them when run on the summary content. + let leaf = seed_leaf( + &cfg, + 0, + "alice@example.com is leading the #launch sprint this week.", + vec![], + vec![], + ); + + let extractor: Arc = Arc::new(CompositeExtractor::regex_only()); + let strategy = LabelStrategy::ExtractFromContent(extractor); + + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &strategy) + .await + .unwrap(); + assert_eq!(sealed.len(), 1, "single 10k-token leaf should seal L0→L1"); + + let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); + assert!( + summary + .entities + .iter() + .any(|e| e == "email:alice@example.com"), + "ExtractFromContent should surface the email entity from summary text; got entities={:?}", + summary.entities + ); + assert!( + summary.topics.iter().any(|t| t == "launch"), + "ExtractFromContent should surface the hashtag-derived topic; got topics={:?}", + summary.topics + ); +} + +#[tokio::test] +async fn seal_with_union_strategy_inherits_labels_from_children() { + let (_tmp, cfg) = test_config(); + let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); + let summariser = InertSummariser::new(); + + // Two leaves with overlapping + distinct labels. Union should + // dedup-merge them into the parent. + let leaf1 = seed_leaf( + &cfg, + 0, + "first leaf body", + vec!["email:alice@example.com".into(), "topic:phoenix".into()], + vec!["phoenix".into(), "launch".into()], + ); + let leaf2 = seed_leaf( + &cfg, + 1, + "second leaf body", + vec!["email:alice@example.com".into(), "person:bob".into()], + vec!["launch".into(), "qa".into()], + ); + + // L0 seals when the budget is crossed. With each leaf at 10k tokens, + // the first append triggers a seal containing only leaf1; we want a + // seal containing both, so use UnionFromChildren and a single seal of + // both leaves at once. The simplest way is to lower budget by sealing + // two leaves into one buffer — the second append crosses budget, so + // the seal contains [leaf1, leaf2]. + // + // Adjust by using smaller token counts so both fit in L0 first, then + // a third append triggers a seal containing both. Reuse the helper + // and override the leaf's token_count for this test. + let leaf1 = LeafRef { + token_count: 5_000, + ..leaf1 + }; + let leaf2 = LeafRef { + token_count: 5_000, + ..leaf2 + }; + + // First leaf: under budget, no seal. + let sealed_1 = append_leaf( + &cfg, + &tree, + &leaf1, + &summariser, + &LabelStrategy::UnionFromChildren, + ) + .await + .unwrap(); + assert!(sealed_1.is_empty()); + // Second leaf: crosses budget → one seal covering both leaves. + let sealed_2 = append_leaf( + &cfg, + &tree, + &leaf2, + &summariser, + &LabelStrategy::UnionFromChildren, + ) + .await + .unwrap(); + assert_eq!(sealed_2.len(), 1); + + let summary = store::get_summary(&cfg, &sealed_2[0]).unwrap().unwrap(); + let entities: std::collections::BTreeSet<&str> = + summary.entities.iter().map(String::as_str).collect(); + let topics: std::collections::BTreeSet<&str> = + summary.topics.iter().map(String::as_str).collect(); + assert!(entities.contains("email:alice@example.com")); + assert!(entities.contains("topic:phoenix")); + assert!(entities.contains("person:bob")); + assert_eq!( + entities.len(), + 3, + "expected 3 unique entities; got {entities:?}" + ); + assert!(topics.contains("phoenix")); + assert!(topics.contains("launch")); + assert!(topics.contains("qa")); + assert_eq!(topics.len(), 3, "expected 3 unique topics; got {topics:?}"); +} + +#[tokio::test] +async fn seal_with_empty_strategy_leaves_labels_empty() { + let (_tmp, cfg) = test_config(); + let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); + let summariser = InertSummariser::new(); + + // Leaf carries labels — Empty strategy should ignore them. + let leaf = seed_leaf( + &cfg, + 0, + "alice@example.com discussing #launch", + vec!["email:alice@example.com".into(), "topic:launch".into()], + vec!["launch".into()], + ); + + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); + assert_eq!(sealed.len(), 1); + + let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); + assert!( + summary.entities.is_empty(), + "Empty strategy must leave entities empty; got {:?}", + summary.entities + ); + assert!( + summary.topics.is_empty(), + "Empty strategy must leave topics empty; got {:?}", + summary.topics + ); +} + +#[tokio::test] +async fn topic_tree_seal_persists_topic_kind_not_source() { + use crate::openhuman::memory::tree::source_tree::types::TreeStatus; + + let (_tmp, cfg) = test_config(); + // Build a topic tree directly — `seal_one_level` runs for both + // source and topic trees, and previously hardcoded Source on the + // resulting summary regardless of the parent tree's kind. + let tree = Tree { + id: "topic-tree-test-id".to_string(), + kind: TreeKind::Topic, + scope: "topic:launch".to_string(), + root_id: None, + max_level: 0, + status: TreeStatus::Active, + created_at: Utc::now(), + last_sealed_at: None, + }; + store::insert_tree(&cfg, &tree).unwrap(); + + let summariser = InertSummariser::new(); + let leaf = seed_leaf(&cfg, 0, "topic content", vec![], vec![]); + + let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) + .await + .unwrap(); + assert_eq!(sealed.len(), 1); + + let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap(); + assert_eq!( + summary.tree_kind, + TreeKind::Topic, + "topic-tree summary must persist tree_kind=Topic, not Source" + ); +} diff --git a/src/openhuman/memory/tree/source_tree/flush.rs b/src/openhuman/memory/tree/source_tree/flush.rs index 920c47061be..0d2710f65bf 100644 --- a/src/openhuman/memory/tree/source_tree/flush.rs +++ b/src/openhuman/memory/tree/source_tree/flush.rs @@ -13,7 +13,7 @@ use anyhow::Result; use chrono::{DateTime, Duration, Utc}; use crate::openhuman::config::Config; -use crate::openhuman::memory::tree::source_tree::bucket_seal::cascade_all_from; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{cascade_all_from, LabelStrategy}; use crate::openhuman::memory::tree::source_tree::store; use crate::openhuman::memory::tree::source_tree::summariser::Summariser; use crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS; @@ -25,6 +25,7 @@ pub async fn flush_stale_buffers( config: &Config, max_age: Duration, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result { let now = Utc::now(); let cutoff = now - max_age; @@ -48,7 +49,8 @@ pub async fn flush_stale_buffers( continue; } }; - let sealed = cascade_all_from(config, &tree, buf.level, summariser, Some(now)).await?; + let sealed = + cascade_all_from(config, &tree, buf.level, summariser, Some(now), strategy).await?; seals += sealed.len(); } Ok(seals) @@ -58,11 +60,13 @@ pub async fn flush_stale_buffers( pub async fn flush_stale_buffers_default( config: &Config, summariser: &dyn Summariser, + strategy: &LabelStrategy, ) -> Result { flush_stale_buffers( config, Duration::seconds(DEFAULT_FLUSH_AGE_SECS), summariser, + strategy, ) .await } @@ -74,10 +78,11 @@ pub async fn force_flush_tree( tree_id: &str, summariser: &dyn Summariser, now: Option>, + strategy: &LabelStrategy, ) -> Result> { let tree = store::get_tree(config, tree_id)? .ok_or_else(|| anyhow::anyhow!("no tree with id {tree_id}"))?; - cascade_all_from(config, &tree, 0, summariser, now).await + cascade_all_from(config, &tree, 0, summariser, now, strategy).await } #[cfg(test)] @@ -136,12 +141,15 @@ mod tests { topics: vec![], score: 0.5, }; - append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); - assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); - - let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser) + append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) .await .unwrap(); + assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); + + let seals = + flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert_eq!(seals, 1); assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 1); @@ -183,11 +191,14 @@ mod tests { topics: vec![], score: 0.5, }; - append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap(); - - let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser) + append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty) .await .unwrap(); + + let seals = + flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty) + .await + .unwrap(); assert_eq!(seals, 0); assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); } diff --git a/src/openhuman/memory/tree/source_tree/mod.rs b/src/openhuman/memory/tree/source_tree/mod.rs index f97d0799590..1f9694468e2 100644 --- a/src/openhuman/memory/tree/source_tree/mod.rs +++ b/src/openhuman/memory/tree/source_tree/mod.rs @@ -22,7 +22,7 @@ pub mod store; pub mod summariser; pub mod types; -pub use bucket_seal::{append_leaf, LeafRef}; +pub use bucket_seal::{append_leaf, append_leaf_deferred, LabelStrategy, LeafRef}; pub use registry::get_or_create_source_tree; pub use store::{get_summary_embedding, set_summary_embedding}; pub use summariser::{build_summariser, inert::InertSummariser, llm::LlmSummariser, Summariser}; diff --git a/src/openhuman/memory/tree/source_tree/summariser/llm.rs b/src/openhuman/memory/tree/source_tree/summariser/llm.rs index 09c8e6ba459..8f9f31135c0 100644 --- a/src/openhuman/memory/tree/source_tree/summariser/llm.rs +++ b/src/openhuman/memory/tree/source_tree/summariser/llm.rs @@ -6,10 +6,12 @@ //! When the source / topic / global tree's bucket-seal cascade decides to //! fold N contributions (raw leaves at L0→L1, or lower-level summaries at //! L_n→L_{n+1}), this summariser is asked to produce the parent node's -//! `content` + derived `entities` + `topics`. The seal machinery itself -//! (bucket budgeting, level promotion, `mem_tree_summaries` persistence) -//! is unchanged — only the text inside the summary row differs from -//! [`super::inert::InertSummariser`]. +//! `content`. The seal machinery itself (bucket budgeting, level +//! promotion, `mem_tree_summaries` persistence) is unchanged — only the +//! text inside the summary row differs from [`super::inert::InertSummariser`]. +//! Entities and topics on `SummaryOutput` are always emitted empty by +//! this summariser; canonical entity ids are populated separately by the +//! entity extractor. //! //! ## Soft-fallback contract //! @@ -18,14 +20,12 @@ //! parent row. We therefore promise **never** to return `Err`: every //! failure (transport, HTTP status, JSON shape) falls back to the same //! deterministic concat-and-truncate behaviour as `InertSummariser` and -//! logs a warn. Callers distinguish "LLM ran fine" from "we fell back" -//! only by observing whether the returned `entities`/`topics` are -//! populated — the inert branch emits empty vecs. +//! logs a warn. //! //! ## Prompt shape //! //! The system prompt commits the model to returning JSON with the shape -//! `{ summary, entities, topics }`. We use Ollama's `format: "json"` + +//! `{ summary }`. We use Ollama's `format: "json"` + //! `temperature: 0.0` to maximise determinism — same knobs the entity //! extractor already uses with success. @@ -40,15 +40,24 @@ use super::inert::InertSummariser; use super::{Summariser, SummaryContext, SummaryInput, SummaryOutput}; use crate::openhuman::memory::tree::types::approx_token_count; -/// Hard cap on summary OUTPUT tokens, regardless of the seal's -/// `ctx.token_budget`. Driven by the embedder's context window — -/// `nomic-embed-text-v1.5` accepts up to 8192 tokens and Phase 4 -/// (`source_tree::bucket_seal`) embeds the summary right after we -/// produce it. If our summary overshoots, the embedder returns 500 -/// and the whole seal transaction rolls back → no summary persists. -/// 6000 leaves a safety margin for tokenizer differences and the -/// JSON wrapper. -const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 6_000; +/// Hard cap on summariser output length (in approximate tokens). +/// +/// Two constraints set this: +/// +/// 1. The downstream embedder (`nomic-embed-text-v1.5`) accepts up to +/// 8192 tokens, and Phase 4 (`source_tree::bucket_seal`) embeds the +/// summary right after we produce it. An overshoot returns HTTP 500 +/// and rolls back the whole seal transaction. +/// 2. Empirically, small instruction-tuned models running locally +/// degrade quickly past ~3500 tokens — they drift, hallucinate, or +/// produce repetitive boilerplate as they extend toward longer +/// targets. Keeping the cap below that breakeven keeps output +/// quality stable on local Ollama deployments. +/// +/// 3500 sits comfortably under the embedder ceiling AND below the local +/// LLM quality cliff. The post-generation [`clamp_to_budget`] enforces +/// this regardless of what the model produces. +const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 3_500; /// Context window we ask Ollama for. Must match the value below in /// [`OllamaOptions::num_ctx`] so the per-input clamp computed in @@ -261,31 +270,18 @@ impl Summariser for LlmSummariser { let (content, token_count) = clamp_to_budget(&parsed.summary, effective_budget); log::debug!( - "[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={} \ - surface_entities_dropped={} topics={}", + "[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={}", ctx.tree_id, ctx.target_level, inputs.len(), - token_count, - parsed.entities.len(), - parsed.topics.len() + token_count ); - // Drop LLM-emitted entities. The model returns surface forms - // ("Alice", "she"), but `SummaryNode.entities` is indexed via - // `index_summary_entity_ids_tx` as canonical ids. Surface forms - // would silently corrupt that index — searches by canonical id - // would not find these summaries. Canonicalisation is the - // entity extractor's job, not the summariser's. Topics stay - // because they're free-form labels, not indexed as canonical - // ids. The `entities` field stays in the prompt to nudge the - // model toward entity-aware summarisation; we just don't - // persist its output. Ok(SummaryOutput { content, token_count, entities: Vec::new(), - topics: dedupe_sorted(parsed.topics), + topics: Vec::new(), }) } } @@ -314,20 +310,21 @@ fn build_user_prompt(inputs: &[SummaryInput], per_input_cap_tokens: u32) -> Stri out } -/// System prompt. Token budget is templated in so the model aims under it. -fn system_prompt(budget: u32) -> String { - format!( - "You are a precise summariser. Summarise the user-provided contributions into a \ - single cohesive passage that preserves concrete facts, decisions, named entities, \ - and temporal ordering. Do not invent facts. Stay well under {budget} tokens.\n\ - \n\ - Return JSON only — no prose, no markdown, no commentary. Schema:\n\ - {{\n\ - \x20 \"summary\": \"\",\n\ - \x20 \"entities\": [\"\"],\n\ - \x20 \"topics\": [\"\"]\n\ - }}" - ) +/// System prompt. Length isn't templated in — empirically, telling small +/// instruction-tuned models "stay under N tokens" makes them produce +/// curt, generic output even when the input has plenty of substance. +/// Output is clamped post-generation by [`clamp_to_budget`] in the +/// caller, so we don't need the model to self-police length. +fn system_prompt(_budget: u32) -> String { + "You are a precise summariser. Summarise the user-provided contributions into a \ + single cohesive passage that preserves concrete facts, decisions, \ + and temporal ordering. Do not invent facts.\n\ + \n\ + Return JSON only — no prose, no markdown, no commentary. Schema:\n\ + {\n\ + \x20 \"summary\": \"\"\n\ + }" + .to_string() } /// Truncate to the caller's token budget using the same ~4 chars/token @@ -343,16 +340,6 @@ fn clamp_to_budget(text: &str, budget: u32) -> (String, u32) { (truncated, tokens) } -fn dedupe_sorted(mut items: Vec) -> Vec { - for item in items.iter_mut() { - *item = item.trim().to_string(); - } - items.retain(|s| !s.is_empty()); - items.sort(); - items.dedup(); - items -} - fn truncate_for_log(s: &str, max_chars: usize) -> String { if s.chars().count() <= max_chars { return s.to_string(); @@ -405,10 +392,6 @@ struct OllamaResponseMessage { struct LlmSummaryOutput { #[serde(default)] summary: String, - #[serde(default)] - entities: Vec, - #[serde(default)] - topics: Vec, } #[cfg(test)] @@ -493,12 +476,16 @@ mod tests { } #[test] - fn system_prompt_templates_budget() { + fn system_prompt_describes_schema() { + // Budget is no longer templated into the prompt — small models + // produced overly curt output when told to "stay under N tokens". + // The clamp in `clamp_to_budget` handles enforcement instead. let p = system_prompt(4096); - assert!(p.contains("4096")); + assert!(!p.contains("4096")); + assert!(!p.contains("Stay well under")); assert!(p.contains("\"summary\"")); - assert!(p.contains("\"entities\"")); - assert!(p.contains("\"topics\"")); + assert!(!p.contains("\"entities\"")); + assert!(!p.contains("\"topics\"")); } #[test] @@ -516,19 +503,6 @@ mod tests { assert!(t <= 6); } - #[test] - fn dedupe_sorted_trims_and_dedupes() { - let out = dedupe_sorted(vec![ - "Bob".into(), - " Alice ".into(), - "Bob".into(), - "".into(), - " ".into(), - "Alice".into(), - ]); - assert_eq!(out, vec!["Alice", "Bob"]); - } - #[test] fn truncate_for_log_short_input_unchanged() { assert_eq!(truncate_for_log("hi", 10), "hi"); @@ -592,16 +566,24 @@ mod tests { assert!(!req.stream); assert_eq!(req.options.temperature, 0.0); assert_eq!(req.messages[0].role, "system"); - assert!(req.messages[0].content.contains("2048")); + assert!(req.messages[0].content.contains("\"summary\"")); assert_eq!(req.messages[1].role, "user"); assert_eq!(req.messages[1].content, "body"); } #[test] - fn llm_output_deserialises_with_missing_fields() { + fn llm_output_deserialises_with_only_summary() { let v: LlmSummaryOutput = serde_json::from_str(r#"{"summary":"hi"}"#).unwrap(); assert_eq!(v.summary, "hi"); - assert!(v.entities.is_empty()); - assert!(v.topics.is_empty()); + } + + #[test] + fn llm_output_ignores_extraneous_fields() { + // Prompt no longer asks for entities/topics, but if the model + // emits them anyway we should still parse `summary` cleanly. + let v: LlmSummaryOutput = + serde_json::from_str(r#"{"summary":"hi","entities":["Alice"],"topics":["x"]}"#) + .unwrap(); + assert_eq!(v.summary, "hi"); } } diff --git a/src/openhuman/memory/tree/store.rs b/src/openhuman/memory/tree/store.rs index 9da6da299bf..1e341beff0c 100644 --- a/src/openhuman/memory/tree/store.rs +++ b/src/openhuman/memory/tree/store.rs @@ -21,6 +21,12 @@ const DEFAULT_LIST_LIMIT: usize = 100; const MAX_LIST_LIMIT: usize = 10_000; const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5); +pub const CHUNK_STATUS_PENDING_EXTRACTION: &str = "pending_extraction"; +pub const CHUNK_STATUS_ADMITTED: &str = "admitted"; +pub const CHUNK_STATUS_BUFFERED: &str = "buffered"; +pub const CHUNK_STATUS_SEALED: &str = "sealed"; +pub const CHUNK_STATUS_DROPPED: &str = "dropped"; + const SCHEMA: &str = " PRAGMA foreign_keys = ON; @@ -174,6 +180,35 @@ CREATE TABLE IF NOT EXISTS mem_tree_entity_hotness ( CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_hotness_score ON mem_tree_entity_hotness(last_hotness); + +-- Async job queue for memory-tree work (extract → admit → buffer → seal → +-- topic-route → daily digest). Producers (ingest, schedulers, handlers) +-- enqueue rows transactionally; the worker pool claims them via the +-- `(status, available_at_ms)` index. `dedupe_key` is enforced as unique +-- only for ready/running rows so a completed job's key can be re-used. +CREATE TABLE IF NOT EXISTS mem_tree_jobs ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL, + payload_json TEXT NOT NULL, + dedupe_key TEXT, + status TEXT NOT NULL DEFAULT 'ready', + attempts INTEGER NOT NULL DEFAULT 0, + max_attempts INTEGER NOT NULL DEFAULT 5, + available_at_ms INTEGER NOT NULL, + locked_until_ms INTEGER, + last_error TEXT, + created_at_ms INTEGER NOT NULL, + started_at_ms INTEGER, + completed_at_ms INTEGER +); + +CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_ready + ON mem_tree_jobs(status, available_at_ms); +CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_kind + ON mem_tree_jobs(kind); +CREATE UNIQUE INDEX IF NOT EXISTS idx_mem_tree_jobs_dedupe_active + ON mem_tree_jobs(dedupe_key) + WHERE dedupe_key IS NOT NULL AND status IN ('ready', 'running'); "; /// Upsert a batch of chunks atomically. @@ -361,6 +396,70 @@ pub fn count_chunks(config: &Config) -> Result { }) } +pub fn set_chunk_lifecycle_status(config: &Config, chunk_id: &str, status: &str) -> Result<()> { + with_connection(config, |conn| { + set_chunk_lifecycle_status_conn(conn, chunk_id, status) + }) +} + +pub(crate) fn set_chunk_lifecycle_status_tx( + tx: &Transaction<'_>, + chunk_id: &str, + status: &str, +) -> Result<()> { + set_chunk_lifecycle_status_conn(tx, chunk_id, status) +} + +pub fn get_chunk_lifecycle_status(config: &Config, chunk_id: &str) -> Result> { + with_connection(config, |conn| { + get_chunk_lifecycle_status_conn(conn, chunk_id) + }) +} + +pub(crate) fn get_chunk_lifecycle_status_tx( + tx: &Transaction<'_>, + chunk_id: &str, +) -> Result> { + get_chunk_lifecycle_status_conn(tx, chunk_id) +} + +fn get_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str) -> Result> { + let row = conn + .query_row( + "SELECT lifecycle_status FROM mem_tree_chunks WHERE id = ?1", + params![chunk_id], + |r| r.get::<_, String>(0), + ) + .optional()?; + Ok(row) +} + +pub fn count_chunks_by_lifecycle_status(config: &Config, status: &str) -> Result { + with_connection(config, |conn| { + let n: i64 = conn.query_row( + "SELECT COUNT(*) FROM mem_tree_chunks WHERE lifecycle_status = ?1", + params![status], + |r| r.get(0), + )?; + Ok(n.max(0) as u64) + }) +} + +fn set_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str, status: &str) -> Result<()> { + let changed = conn.execute( + "UPDATE mem_tree_chunks SET lifecycle_status = ?1 WHERE id = ?2", + params![status, chunk_id], + )?; + if changed == 0 { + log::warn!( + "[memory_tree::store] lifecycle update affected 0 rows chunk_id={} status={}", + chunk_id, + status + ); + } + Ok(()) +} + fn row_to_chunk(row: &rusqlite::Row<'_>) -> rusqlite::Result { let id: String = row.get(0)?; let source_kind_s: String = row.get(1)?; @@ -451,6 +550,21 @@ pub(crate) fn with_connection( // legacy summaries from Phases 1-3 read back as None; retrieval // tolerates NULL by dropping the row to the bottom of a rerank. add_column_if_missing(&conn, "mem_tree_summaries", "embedding", "BLOB")?; + // Async-pipeline lifecycle flag. Default 'admitted' so chunks ingested + // before the queue migration stay queryable. New writes start at + // 'pending_extraction'; the extract handler advances them to 'admitted' + // (then 'buffered' / 'sealed') or 'dropped'. + add_column_if_missing( + &conn, + "mem_tree_chunks", + "lifecycle_status", + "TEXT NOT NULL DEFAULT 'admitted'", + )?; + conn.execute_batch( + "CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_lifecycle \ + ON mem_tree_chunks(lifecycle_status);", + ) + .context("Failed to create mem_tree_chunks lifecycle index")?; f(&conn) } diff --git a/src/openhuman/memory/tree/topic_tree/backfill.rs b/src/openhuman/memory/tree/topic_tree/backfill.rs index e98a53e532e..6d0518df54b 100644 --- a/src/openhuman/memory/tree/topic_tree/backfill.rs +++ b/src/openhuman/memory/tree/topic_tree/backfill.rs @@ -1,22 +1,36 @@ //! Topic-tree backfill — hydrate a freshly-materialised topic tree with -//! every historical leaf mentioning the entity (#709 Phase 3c). +//! recent leaves mentioning the entity (#709 Phase 3c). //! //! When the curator decides an entity has crossed the hotness threshold //! for the first time, we create a fresh topic tree AND walk the -//! `mem_tree_entity_index` inverted index to append every prior leaf into +//! `mem_tree_entity_index` inverted index to append matching leaves into //! its L0 buffer. Reusing `bucket_seal::append_leaf` means the cascade -//! fires automatically — a well-established entity may seal several -//! levels as soon as the tree is spawned. +//! fires automatically. +//! +//! ## Why bounded by hotness window +//! +//! Hotness uses a 30-day recency decay (see `topic_tree::hotness`). Leaves +//! older than 30 days contribute zero to current hotness, so by definition +//! they cannot be the reason a tree is spawning *now*. Including them +//! bloats the spawn latency, wastes summariser LLM calls, and amplifies +//! ancient signal that has already decayed away. We cap the backfill +//! window at [`BACKFILL_WINDOW_DAYS`] to align with the hotness math. +//! +//! Older content is still queryable through source-tree retrieval and the +//! entity index — it just doesn't get its own slot in the topic tree. //! //! Backfill is intentionally best-effort: missing chunks are skipped with //! a warn log rather than failing the whole spawn, because Phase 3c is //! additive — a partial topic tree is still useful. use anyhow::{Context, Result}; +use chrono::Utc; use crate::openhuman::config::Config; use crate::openhuman::memory::tree::score::store::lookup_entity; -use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, +}; use crate::openhuman::memory::tree::source_tree::summariser::Summariser; use crate::openhuman::memory::tree::source_tree::types::Tree; use crate::openhuman::memory::tree::store::get_chunk; @@ -25,6 +39,13 @@ use crate::openhuman::memory::tree::store::get_chunk; /// keeps initial spawn latency bounded even for very active entities. const BACKFILL_LIMIT: usize = 500; +/// Backfill window in days — matches `topic_tree::hotness::recency_decay`'s +/// hard cliff. Leaves older than this contribute zero to current hotness +/// so they cannot have driven the spawn decision. +pub const BACKFILL_WINDOW_DAYS: i64 = 30; + +const DAY_MS: i64 = 24 * 60 * 60 * 1_000; + /// Walk the entity index for `entity_id` and append every discovered leaf /// to `tree`. Returns the number of leaves appended (NOT the number of /// summaries sealed). Idempotent: `append_leaf` itself is a no-op when a @@ -35,10 +56,33 @@ pub async fn backfill_topic_tree( entity_id: &str, summariser: &dyn Summariser, ) -> Result { + backfill_topic_tree_at( + config, + tree, + entity_id, + summariser, + Utc::now().timestamp_millis(), + ) + .await +} + +/// Deterministic variant — backfill against a caller-supplied `now_ms` +/// for the recency window. Used by tests so the 30-day cutoff doesn't +/// depend on the wall clock. +pub async fn backfill_topic_tree_at( + config: &Config, + tree: &Tree, + entity_id: &str, + summariser: &dyn Summariser, + now_ms: i64, +) -> Result { + let cutoff_ms = now_ms.saturating_sub(BACKFILL_WINDOW_DAYS.saturating_mul(DAY_MS)); log::info!( - "[topic_tree::backfill] start entity_id={} tree_id={}", + "[topic_tree::backfill] start entity_id={} tree_id={} window_days={} cutoff_ms={}", entity_id, - tree.id + tree.id, + BACKFILL_WINDOW_DAYS, + cutoff_ms ); let hits = lookup_entity(config, entity_id, Some(BACKFILL_LIMIT)) @@ -52,10 +96,30 @@ pub async fn backfill_topic_tree( return Ok(0); } + // Drop hits older than the hotness recency window — see module docs. + let total_hits = hits.len(); + let mut hits: Vec<_> = hits + .into_iter() + .filter(|h| h.timestamp_ms >= cutoff_ms) + .collect(); + let dropped = total_hits - hits.len(); + if dropped > 0 { + log::debug!( + "[topic_tree::backfill] dropped {dropped} hits older than {BACKFILL_WINDOW_DAYS}d \ + for entity_id={entity_id}" + ); + } + if hits.is_empty() { + log::debug!( + "[topic_tree::backfill] all entity-index hits fell outside the {BACKFILL_WINDOW_DAYS}d \ + window for entity_id={entity_id} — empty backfill" + ); + return Ok(0); + } + // Sort by timestamp ASC so the buffer's `oldest_at` and the sealed // summary's `time_range_start` reflect the true historical order, not // the DESC ordering `lookup_entity` returns. - let mut hits = hits; hits.sort_by_key(|h| h.timestamp_ms); let mut appended = 0usize; @@ -95,7 +159,10 @@ pub async fn backfill_topic_tree( score: hit.score, }; - append_leaf(config, tree, &leaf, summariser) + // Topic-tree backfill: empty labels for sealed summaries — the + // tree's scope already pins the canonical id, so cross-pollinating + // descendants' entities would noise the index. See LabelStrategy. + append_leaf(config, tree, &leaf, summariser, &LabelStrategy::Empty) .await .with_context(|| { format!( @@ -172,6 +239,11 @@ mod tests { } } + /// Deterministic "now" used by the windowed-backfill tests: 1 hour + /// after the latest seeded leaf so all three sit inside the 30-day + /// cutoff. Lets us keep the legacy 2023-era timestamps unchanged. + const TEST_NOW_MS: i64 = 1_700_000_020_000 + 3_600_000; + #[tokio::test] async fn backfill_appends_all_entity_leaves() { let (_tmp, cfg) = test_config(); @@ -212,9 +284,15 @@ mod tests { let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); assert_eq!(n, 3); // L0 buffer should hold all three leaves (combined tokens well @@ -226,6 +304,38 @@ mod tests { assert_eq!(buf.oldest_at.unwrap().timestamp_millis(), 1_700_000_000_000); } + #[tokio::test] + async fn backfill_drops_leaves_older_than_window() { + let (_tmp, cfg) = test_config(); + // c_old is 60d before TEST_NOW_MS — outside the 30d cutoff. + // c_new is 5d before TEST_NOW_MS — inside the window. + let old_ts = TEST_NOW_MS - 60 * DAY_MS; + let new_ts = TEST_NOW_MS - 5 * DAY_MS; + let c_old = mk_chunk("slack:#eng", 0, old_ts, 100); + let c_new = mk_chunk("slack:#eng", 1, new_ts, 100); + upsert_chunks(&cfg, &[c_old.clone(), c_new.clone()]).unwrap(); + + let e = sample_entity("email:alice@example.com", "alice@example.com"); + index_entity(&cfg, &e, &c_old.id, "leaf", old_ts, Some("source:slack")).unwrap(); + index_entity(&cfg, &e, &c_new.id, "leaf", new_ts, Some("source:slack")).unwrap(); + + let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); + let summariser = InertSummariser::new(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); + assert_eq!(n, 1, "only the in-window leaf should be appended"); + let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap(); + assert_eq!(buf.item_ids.len(), 1); + assert_eq!(buf.item_ids[0], c_new.id); + } + #[tokio::test] async fn backfill_skips_missing_chunks_without_failing() { let (_tmp, cfg) = test_config(); @@ -247,9 +357,15 @@ mod tests { let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); assert_eq!(n, 1, "only the existing chunk should be appended"); let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap(); assert_eq!(buf.item_ids.len(), 1); @@ -273,12 +389,24 @@ mod tests { let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); - backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); + backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); // append_leaf is idempotent so the buffer still has exactly one row. let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap(); assert_eq!(buf.item_ids.len(), 1); @@ -300,9 +428,15 @@ mod tests { .unwrap(); let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap(); let summariser = InertSummariser::new(); - let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser) - .await - .unwrap(); + let n = backfill_topic_tree_at( + &cfg, + &tree, + "email:alice@example.com", + &summariser, + TEST_NOW_MS, + ) + .await + .unwrap(); assert_eq!(n, 0); } } diff --git a/src/openhuman/memory/tree/topic_tree/curator.rs b/src/openhuman/memory/tree/topic_tree/curator.rs index ae4006272b1..a911697a220 100644 --- a/src/openhuman/memory/tree/topic_tree/curator.rs +++ b/src/openhuman/memory/tree/topic_tree/curator.rs @@ -186,7 +186,11 @@ mod tests { } fn seed_leaf_for_entity(cfg: &Config, entity_id: &str, source_tree: &str, seq: u32) { - let ts_ms = 1_700_000_000_000 + (seq as i64) * 1_000; + // Use a "now-anchored" timestamp so backfill's 30-day window + // (see topic_tree::backfill::BACKFILL_WINDOW_DAYS) always + // includes these seeded leaves. Spread by seq to keep ordering + // deterministic. + let ts_ms = Utc::now().timestamp_millis() - (seq as i64) * 1_000; let ts = Utc.timestamp_millis_opt(ts_ms).unwrap(); let c = Chunk { id: chunk_id(SourceKind::Chat, source_tree, seq, "test-content"), diff --git a/src/openhuman/memory/tree/topic_tree/routing.rs b/src/openhuman/memory/tree/topic_tree/routing.rs index 6867561255b..7d315353be2 100644 --- a/src/openhuman/memory/tree/topic_tree/routing.rs +++ b/src/openhuman/memory/tree/topic_tree/routing.rs @@ -21,7 +21,9 @@ use anyhow::Result; use crate::openhuman::config::Config; -use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef}; +use crate::openhuman::memory::tree::source_tree::bucket_seal::{ + append_leaf, LabelStrategy, LeafRef, +}; use crate::openhuman::memory::tree::source_tree::store as src_store; use crate::openhuman::memory::tree::source_tree::summariser::Summariser; use crate::openhuman::memory::tree::source_tree::types::{TreeKind, TreeStatus}; @@ -90,7 +92,16 @@ async fn route_one_entity( entities: vec![entity_id.to_string()], ..leaf.clone() }; - append_leaf(config, &tree, &topic_leaf, summariser).await?; + // Topic-tree seals leave entities/topics empty: the tree's + // scope already pins the canonical id this tree represents. + append_leaf( + config, + &tree, + &topic_leaf, + summariser, + &LabelStrategy::Empty, + ) + .await?; } else { log::debug!( "[topic_tree::routing] skip archived topic tree id={} entity={}", @@ -295,9 +306,15 @@ mod tests { crate::openhuman::memory::tree::topic_tree::types::TOPIC_RECHECK_EVERY - 1; crate::openhuman::memory::tree::topic_tree::store::upsert(&cfg, &counters).unwrap(); - // Seed a leaf in slack and gmail referencing Alice. - let c1 = persist_chunk(&cfg, "slack:#eng", 0, 1_700_000_000_000, 100); - let c2 = persist_chunk(&cfg, "gmail:alice", 0, 1_700_000_010_000, 100); + // Seed leaves in slack and gmail referencing Alice. Anchor the + // timestamps to "now" so the 30-day backfill window + // (topic_tree::backfill::BACKFILL_WINDOW_DAYS) covers them. + let now_ms = Utc::now().timestamp_millis(); + let ts_c1 = now_ms - 20_000; + let ts_c2 = now_ms - 10_000; + let ts_c3 = now_ms; + let c1 = persist_chunk(&cfg, "slack:#eng", 0, ts_c1, 100); + let c2 = persist_chunk(&cfg, "gmail:alice", 0, ts_c2, 100); let e = CanonicalEntity { canonical_id: entity_id.into(), kind: EntityKind::Email, @@ -306,24 +323,16 @@ mod tests { span_end: entity_id.len() as u32, score: 1.0, }; - index_entity(&cfg, &e, &c1, "leaf", 1_700_000_000_000, Some("slack:#eng")).unwrap(); - index_entity( - &cfg, - &e, - &c2, - "leaf", - 1_700_000_010_000, - Some("gmail:alice"), - ) - .unwrap(); + index_entity(&cfg, &e, &c1, "leaf", ts_c1, Some("slack:#eng")).unwrap(); + index_entity(&cfg, &e, &c2, "leaf", ts_c2, Some("gmail:alice")).unwrap(); // A third leaf arrives — should both fan out to (future) topic tree // and push the curator over the recheck cadence, materialising it. - let c3 = persist_chunk(&cfg, "slack:#eng", 1, 1_700_000_020_000, 100); + let c3 = persist_chunk(&cfg, "slack:#eng", 1, ts_c3, 100); let leaf = LeafRef { chunk_id: c3.clone(), token_count: 100, - timestamp: Utc.timestamp_millis_opt(1_700_000_020_000).unwrap(), + timestamp: Utc.timestamp_millis_opt(ts_c3).unwrap(), content: "new mention".into(), entities: vec![entity_id.into()], topics: vec![], diff --git a/tests/json_rpc_e2e.rs b/tests/json_rpc_e2e.rs index eb15b8794e1..54cd8c27a26 100644 --- a/tests/json_rpc_e2e.rs +++ b/tests/json_rpc_e2e.rs @@ -825,6 +825,7 @@ async fn json_rpc_memory_tree_end_to_end() { "openhuman.memory_tree_ingest".to_string(), "openhuman.memory_tree_list_chunks".to_string(), "openhuman.memory_tree_get_chunk".to_string(), + "openhuman.memory_tree_trigger_digest".to_string(), ]; assert_eq!(controllers.len(), expected_methods.len()); for method in &expected_methods {