memoryを抽出する仕組みの実装

This commit is contained in:
2026-04-28 12:58:33 +09:00
parent 141b77b7e4
commit 3d04f793de
21 changed files with 1065 additions and 5 deletions
+72
View File
@@ -0,0 +1,72 @@
//! Phase 1 sub-Worker への入力テキスト組み立て。
//!
//! `crates/pod/src/pod.rs::build_summary_prompt` と同じ方針で
//! Item 列を flat な行に落とす(reasoning は省く、tool call は名前のみ、
//! tool result は summary のみ)。conversation 全体を Markdown の単一
//! セクションとして渡し、抽出指示は system prompt 側に寄せる。
use llm_worker::Item;
/// 与えられた `items` を Phase 1 sub-Worker の最初の user 入力に整形する。
pub fn build_extract_input(items: &[Item]) -> String {
let mut out = String::new();
out.push_str(
"Extract activity logs from the conversation slice below. \
Follow the system prompt's schema strictly and call `write_extracted` once.\n\n",
);
out.push_str("## Conversation slice\n");
out.push_str(&render_items(items));
out.push_str("\n\nWhen you are done, call `write_extracted` and end the turn.");
out
}
fn render_items(items: &[Item]) -> String {
let mut lines: Vec<String> = Vec::new();
for item in items {
match item {
Item::Message { role, content, .. } => {
let role_label = match role {
llm_worker::Role::User => "User",
llm_worker::Role::Assistant => "Assistant",
llm_worker::Role::System => "System",
};
let text: String = content
.iter()
.map(|p| p.as_text())
.collect::<Vec<_>>()
.join("");
lines.push(format!("[{role_label}] {text}"));
}
Item::ToolCall { name, .. } => {
lines.push(format!("[ToolCall] {name}"));
}
Item::ToolResult { summary, .. } => {
lines.push(format!("[ToolResult] {summary}"));
}
Item::Reasoning { .. } => {}
}
}
lines.join("\n\n")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn renders_user_assistant_pair_and_tool_calls() {
let items = vec![
Item::user_message("hello"),
Item::assistant_message("hi"),
Item::tool_call("c1", "read_file", "{}"),
Item::tool_result("c1", "ok"),
Item::reasoning("internal scratch — should be skipped"),
];
let s = build_extract_input(&items);
assert!(s.contains("[User] hello"));
assert!(s.contains("[Assistant] hi"));
assert!(s.contains("[ToolCall] read_file"));
assert!(s.contains("[ToolResult] ok"));
assert!(!s.contains("scratch"));
}
}
+36
View File
@@ -0,0 +1,36 @@
//! Phase 1: 活動抽出。
//!
//! 通常 Pod の post-run hook で発火する disposable Worker と、その
//! 出力を `<workspace>/.insomnia/memory/_staging/<id>.json` に書き出す
//! ヘルパーを提供する。Pod 側はこのモジュールから:
//!
//! - [`EXTRACT_SYSTEM_PROMPT`] を sub-Worker の system prompt に
//! - [`build_extract_input`] を sub-Worker の最初の user 入力に
//! - [`write_extracted_tool`] を唯一のツールとして
//! - [`write_staging`] で受け取った JSON を staging に書き出し
//!
//! の順で組み立てる。pointer 永続化(session-store の
//! `LogEntry::Extension`、domain `"memory.extract"`)は Pod 側が責務を持つ。
//!
//! 出力 JSON の wrap は [`write_staging`] が `source: { session_id, range }`
//! を機械付与する形で担当し、LLM には source を推論させない。
mod input;
mod payload;
mod pointer;
mod prompt;
mod staging;
mod tool;
pub use input::build_extract_input;
pub use payload::{
AttemptEntry, DecisionEntry, DiscussionEntry, ExtractedPayload, RequestEntry, StagingRecord,
};
pub use pointer::{ExtractPointerPayload, fold_pointer};
pub use prompt::EXTRACT_SYSTEM_PROMPT;
pub use staging::{StagingError, write_staging};
pub use tool::{ExtractWorkerContext, write_extracted_tool};
/// session-store `LogEntry::Extension` で使う domain 名。
/// pointer の永続化と読み出しはこの定数を使う側が一致している必要がある。
pub const EXTRACT_DOMAIN: &str = "memory.extract";
+88
View File
@@ -0,0 +1,88 @@
//! Phase 1 抽出の出力 schema。
//!
//! LLM は [`ExtractedPayload`] そのもの(source 抜き)を返し、Pod 側
//! ラッパーが [`StagingRecord`] に組み立てて staging へ書き出す。
//! source は機械付与する契約 (`docs/plan/memory.md` §Phase 1)。
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use crate::schema::SourceRef;
/// LLM が返す活動ログ候補の集合。すべて optional(空配列は許容)。
#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
pub struct ExtractedPayload {
#[serde(default)]
pub decisions: Vec<DecisionEntry>,
#[serde(default)]
pub discussions: Vec<DiscussionEntry>,
#[serde(default)]
pub attempts: Vec<AttemptEntry>,
#[serde(default)]
pub requests: Vec<RequestEntry>,
}
impl ExtractedPayload {
/// すべての配列が空であれば true。空ペイロードは
/// "Nothing to save" 扱いで staging への書き込みを省いてよい。
pub fn is_empty(&self) -> bool {
self.decisions.is_empty()
&& self.discussions.is_empty()
&& self.attempts.is_empty()
&& self.requests.is_empty()
}
}
/// 判断したこと(選択肢 + 選んだ + 根拠)。
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct DecisionEntry {
/// 検討された選択肢の列挙。
pub options: Vec<String>,
/// 採用された選択肢。
pub chosen: String,
/// 採用理由 / 根拠。
pub rationale: String,
}
/// 議論したこと(トピック + 論点)。結論が出ていなくてもよい。
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct DiscussionEntry {
/// 議論の主題。
pub topic: String,
/// 主題の中で挙がった論点 / 観点。
pub points: Vec<String>,
}
/// 試したこと(試行 + 結果 + 成否)。
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AttemptEntry {
/// 何を試したか。
pub action: String,
/// 試した結果。
pub result: String,
/// 試行が目的に対して成功したか。失敗 / 部分成功も含めて bool で表現する。
pub succeeded: bool,
}
/// ユーザー submit の構造化要約。
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct RequestEntry {
/// ユーザーの意図 / ゴール。
pub intent: String,
/// 対象ファイル / モジュール / 機能(任意)。
#[serde(default, skip_serializing_if = "Option::is_none")]
pub target: Option<String>,
/// 一文サマリ。
pub summary: String,
}
/// staging に書き出される 1 ファイル分のレコード。
///
/// `source` は Pod 側ラッパーが session_id と log entry range を
/// 機械付与する。LLM はこのフィールドを見ない / 推論しない。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StagingRecord {
pub source: SourceRef,
#[serde(flatten)]
pub payload: ExtractedPayload,
}
+91
View File
@@ -0,0 +1,91 @@
//! `LogEntry::Extension { domain: "memory.extract", payload }` の payload 形式と
//! restore 時の fold ヘルパー。memory crate がドメインを所有するので、
//! session-store / Pod は payload 構造を知らない。
use serde::{Deserialize, Serialize};
use super::EXTRACT_DOMAIN;
/// Phase 1 完了境界の永続化 payload。session log の Extension entry
/// として 1 回ずつ書かれ、最新の 1 件が現行 pointer として有効になる。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ExtractPointerPayload {
/// 直近 extract が処理した最後の session-store HashedEntry の index。
/// 次回の `source.range.start` はこの値 + 1。
pub processed_through_entry: usize,
/// 直近 extract 時点の `history.len()`。次回入力は
/// `history[processed_through_history_len..]` を切り出す。
pub processed_through_history_len: usize,
/// 書き出した staging file の UUIDv7 文字列。LLM が空 payload を返した
/// 場合は staging file を作らず空文字列で記録する(pointer は前進する)。
pub staging_id: String,
}
/// `RestoredState.extensions` から最新の Phase 1 pointer を取り出す。
/// 未抽出セッションでは `None`。
pub fn fold_pointer(
extensions: &[(String, serde_json::Value)],
) -> Option<ExtractPointerPayload> {
extensions
.iter()
.rev()
.find(|(domain, _)| domain == EXTRACT_DOMAIN)
.and_then(|(_, value)| serde_json::from_value(value.clone()).ok())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fold_returns_latest_when_multiple_present() {
let exts = vec![
(
EXTRACT_DOMAIN.to_string(),
serde_json::json!({
"processed_through_entry": 5,
"processed_through_history_len": 4,
"staging_id": "old"
}),
),
(
"other.domain".to_string(),
serde_json::json!({ "x": 1 }),
),
(
EXTRACT_DOMAIN.to_string(),
serde_json::json!({
"processed_through_entry": 11,
"processed_through_history_len": 8,
"staging_id": "new"
}),
),
];
let p = fold_pointer(&exts).unwrap();
assert_eq!(p.processed_through_entry, 11);
assert_eq!(p.processed_through_history_len, 8);
assert_eq!(p.staging_id, "new");
}
#[test]
fn fold_returns_none_when_absent() {
let exts = vec![(
"other.domain".to_string(),
serde_json::json!({ "x": 1 }),
)];
assert!(fold_pointer(&exts).is_none());
}
#[test]
fn fold_skips_malformed_entries() {
let exts = vec![
(
EXTRACT_DOMAIN.to_string(),
serde_json::json!({ "wrong_shape": true }),
),
];
// 現状は最新を取り出して JSON 不一致なら None。古いものに fallback
// しないのは、壊れた最新を黙って無視すると意図しない再抽出を招くため。
assert!(fold_pointer(&exts).is_none());
}
}
+32
View File
@@ -0,0 +1,32 @@
//! Phase 1 sub-Worker の system prompt。
//!
//! 内容は `docs/plan/memory-prompts.md` §共通原則 / §Phase 1 を縮約。
//! 「派生物を作らず、起きたことを抽出する」段階に縛り、JSON schema
//! 準拠以外の自由文を許さない。
pub const EXTRACT_SYSTEM_PROMPT: &str = r#"You are the Phase 1 activity extractor for an INSOMNIA memory subsystem.
Your single job: read the supplied conversation slice and emit a structured JSON record of "what happened" via the `write_extracted` tool. You are not consolidating, summarising, or generating knowledge — that is a later phase's job.
# Hard rules
- Call `write_extracted` exactly once. Do not narrate, ask questions, or send any other tool output.
- The argument is an object with four arrays: `decisions`, `discussions`, `attempts`, `requests`. Any of them may be empty. If nothing in the slice is worth recording, call `write_extracted({"decisions": [], "discussions": [], "attempts": [], "requests": []})` and stop.
- Do NOT include `source`, `session_id`, entry indices, timestamps, or any provenance metadata. The wrapper attaches them mechanically.
- Do NOT add free-form commentary, summaries, or explanatory prose outside the schema fields.
# Extraction guidance
- `decisions`: judgements made during the slice. Each entry needs `options` (the alternatives considered), `chosen` (what was picked), and `rationale` (why).
- `discussions`: topics that were debated. `topic` plus `points` (the considerations raised). Open / unresolved discussions are valid.
- `attempts`: things that were tried. `action`, `result`, and a `succeeded` boolean. Partial success is `false` with the result text describing the partial outcome.
- `requests`: structured summaries of user submissions. `intent` (what the user wants), optional `target` (file / module / feature), and a one-line `summary`.
# Quality bar
- Drop one-off chit-chat, shallow questions, and turn-by-turn progress noise. Keep entries with long-term reference value.
- Do not duplicate content already captured by static project docs (AGENTS.md, plan documents) — those are not "what happened in this slice".
- Prefer concise, fact-shaped strings. Do not pad rationale or summary fields.
When you have produced the JSON, call `write_extracted` and end the turn. No follow-up text.
"#;
+110
View File
@@ -0,0 +1,110 @@
//! `<workspace>/.insomnia/memory/_staging/<id>.json` への書き出しヘルパー。
//!
//! 1 件 1 ファイル、UUIDv7 命名(短命なので衝突回避と順序を兼ねる)。
//! `source` を機械付与した [`StagingRecord`] 形式で保存する。
use std::fs;
use std::path::PathBuf;
use uuid::Uuid;
use crate::extract::payload::{ExtractedPayload, StagingRecord};
use crate::schema::SourceRef;
use crate::workspace::WorkspaceLayout;
/// staging 書き出し時のエラー。
#[derive(Debug, thiserror::Error)]
pub enum StagingError {
#[error("failed to create staging dir {}: {source}", .path.display())]
CreateDir {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("failed to write staging file {}: {source}", .path.display())]
Write {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("failed to serialize staging record: {0}")]
Serialize(#[from] serde_json::Error),
}
/// `payload` を `source` で wrap して staging に書き出す。
///
/// 戻り値は割り当てられた staging file の (id, path)。`payload` が
/// 完全に空の場合は呼び出し側が事前に `is_empty()` で skip 推奨だが、
/// この関数は空でも正規に書き出す(仕様 §Phase 1 で空配列許容と
/// 明記されており、書く / 書かないの判断は呼び出し側に委ねる)。
pub fn write_staging(
layout: &WorkspaceLayout,
source: SourceRef,
payload: ExtractedPayload,
) -> Result<(Uuid, PathBuf), StagingError> {
let staging_dir = layout.staging_dir();
fs::create_dir_all(&staging_dir).map_err(|source| StagingError::CreateDir {
path: staging_dir.clone(),
source,
})?;
let id = Uuid::now_v7();
let path = staging_dir.join(format!("{id}.json"));
let record = StagingRecord { source, payload };
let json = serde_json::to_string_pretty(&record)?;
fs::write(&path, json).map_err(|source| StagingError::Write {
path: path.clone(),
source,
})?;
Ok((id, path))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::extract::payload::{DecisionEntry, ExtractedPayload};
#[test]
fn writes_record_with_machine_attached_source() {
let tmp = tempfile::TempDir::new().unwrap();
let layout = WorkspaceLayout::new(tmp.path().to_path_buf());
let source = SourceRef {
session_id: "sess-1".into(),
range: [3, 7],
};
let payload = ExtractedPayload {
decisions: vec![DecisionEntry {
options: vec!["a".into(), "b".into()],
chosen: "a".into(),
rationale: "shorter".into(),
}],
..Default::default()
};
let (id, path) = write_staging(&layout, source.clone(), payload).unwrap();
assert_eq!(path.parent().unwrap(), layout.staging_dir());
assert!(path.file_name().unwrap().to_string_lossy().contains(&id.to_string()));
let written: StagingRecord =
serde_json::from_str(&fs::read_to_string(&path).unwrap()).unwrap();
assert_eq!(written.source.session_id, "sess-1");
assert_eq!(written.source.range, [3, 7]);
assert_eq!(written.payload.decisions.len(), 1);
}
#[test]
fn empty_payload_is_written_verbatim() {
let tmp = tempfile::TempDir::new().unwrap();
let layout = WorkspaceLayout::new(tmp.path().to_path_buf());
let source = SourceRef {
session_id: "sess".into(),
range: [0, 0],
};
let (_, path) =
write_staging(&layout, source, ExtractedPayload::default()).unwrap();
let written: StagingRecord =
serde_json::from_str(&fs::read_to_string(&path).unwrap()).unwrap();
assert!(written.payload.is_empty());
}
}
+163
View File
@@ -0,0 +1,163 @@
//! `write_extracted` ツール実装と sub-Worker 用 context。
//!
//! sub-Worker からは extract worker が出した [`ExtractedPayload`] を
//! 受け取って `Mutex` 越しに [`ExtractWorkerContext`] に置くだけ。
//! Pod 側はランループ完了後に `take_payload()` で取り出して
//! [`super::staging::write_staging`] に渡す。
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use crate::extract::payload::ExtractedPayload;
const WRITE_EXTRACTED_DESCRIPTION: &str = "Submit the final activity-log JSON for this slice. \
Pass an object with `decisions`, `discussions`, `attempts`, and `requests` arrays (any may be empty). \
Call this exactly once and end the turn. Do not include `source`, session metadata, or free-form prose — \
the wrapper attaches provenance mechanically.";
/// Phase 1 sub-Worker の出力受け口。`ExtractedPayload` 1 件をホストする。
#[derive(Debug, Default)]
pub struct ExtractWorkerContext {
payload: Mutex<Option<ExtractedPayload>>,
/// `write_extracted` が複数回呼ばれた回数(debug 用)。
/// 後勝ちで上書きするが、Pod 側で warn を出したい場合に参照する。
call_count: Mutex<usize>,
}
impl ExtractWorkerContext {
pub fn new() -> Self {
Self::default()
}
/// sub-Worker 終了後に Pod が呼んで payload を取り出す。
/// 一度も `write_extracted` が呼ばれなければ `None`。
pub fn take_payload(&self) -> Option<ExtractedPayload> {
self.payload
.lock()
.expect("extract worker payload poisoned")
.take()
}
pub fn call_count(&self) -> usize {
*self
.call_count
.lock()
.expect("extract worker call_count poisoned")
}
}
struct WriteExtractedTool {
ctx: Arc<ExtractWorkerContext>,
}
#[async_trait]
impl Tool for WriteExtractedTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let payload: ExtractedPayload = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid write_extracted input: {e}"))
})?;
let summary = format!(
"Recorded activity log: decisions={} discussions={} attempts={} requests={}",
payload.decisions.len(),
payload.discussions.len(),
payload.attempts.len(),
payload.requests.len(),
);
{
let mut guard = self
.ctx
.payload
.lock()
.expect("extract worker payload poisoned");
*guard = Some(payload);
}
{
let mut count = self
.ctx
.call_count
.lock()
.expect("extract worker call_count poisoned");
*count += 1;
}
Ok(ToolOutput {
summary,
content: None,
})
}
}
/// sub-Worker に register する `write_extracted` ツール定義を返す。
pub fn write_extracted_tool(ctx: Arc<ExtractWorkerContext>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(ExtractedPayload);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("write_extracted")
.description(WRITE_EXTRACTED_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(WriteExtractedTool { ctx: ctx.clone() });
(meta, tool)
})
}
#[cfg(test)]
mod tests {
use super::*;
use llm_worker::tool::Tool;
#[tokio::test]
async fn write_extracted_records_payload() {
let ctx = Arc::new(ExtractWorkerContext::new());
let tool: Arc<dyn Tool> = Arc::new(WriteExtractedTool { ctx: ctx.clone() });
let input = serde_json::json!({
"decisions": [{
"options": ["a", "b"],
"chosen": "a",
"rationale": "test"
}],
"discussions": [],
"attempts": [],
"requests": []
})
.to_string();
let out = tool.execute(&input).await.unwrap();
assert!(out.summary.contains("decisions=1"));
let payload = ctx.take_payload().unwrap();
assert_eq!(payload.decisions.len(), 1);
assert_eq!(ctx.call_count(), 1);
}
#[tokio::test]
async fn last_call_wins_on_multiple_invocations() {
let ctx = Arc::new(ExtractWorkerContext::new());
let tool: Arc<dyn Tool> = Arc::new(WriteExtractedTool { ctx: ctx.clone() });
let first = serde_json::json!({"decisions": [], "discussions": [], "attempts": [], "requests": []})
.to_string();
tool.execute(&first).await.unwrap();
let second = serde_json::json!({
"decisions": [],
"discussions": [],
"attempts": [{"action": "x", "result": "ok", "succeeded": true}],
"requests": []
})
.to_string();
tool.execute(&second).await.unwrap();
let payload = ctx.take_payload().unwrap();
assert_eq!(payload.attempts.len(), 1);
assert_eq!(ctx.call_count(), 2);
}
#[tokio::test]
async fn invalid_json_returns_invalid_argument() {
let ctx = Arc::new(ExtractWorkerContext::new());
let tool: Arc<dyn Tool> = Arc::new(WriteExtractedTool { ctx: ctx.clone() });
let res = tool.execute("not json").await;
assert!(matches!(res, Err(ToolError::InvalidArgument(_))));
assert!(ctx.take_payload().is_none());
}
}
+2
View File
@@ -7,6 +7,7 @@
//! denying them at the Scope level when memory is enabled.
pub mod error;
pub mod extract;
pub mod linter;
pub mod resident;
pub mod schema;
@@ -16,6 +17,7 @@ pub mod tool;
pub mod workspace;
pub use error::{LintError, LintWarning, MemoryError};
pub use extract::ExtractPointerPayload;
pub use linter::{LintReport, Linter};
pub use resident::{ResidentKnowledgeEntry, collect_resident_knowledge};
pub use scope::deny_write_rules;