Skip to main content

somatize_mcp/tools/
knowledge.rs

1//! The experiment-pool tools — Case-Based Reasoning over past runs.
2//!
3//! The four CBR steps map onto the tools directly: **Retrieve** is
4//! [`find_similar`], **Reuse** is reading the `run_dir` a hit points at,
5//! **Revise** is [`branch_from`] followed by an actual run, and
6//! **Retain** is [`record_conclusion`].
7//!
8//! Handlers stay thin on purpose: they parse arguments, ask the
9//! knowledge base, and hand the answer to [`crate::render`]. All the
10//! text a model sees is produced by pure functions over there, where it
11//! is snapshot-tested.
12
13use crate::context::SomaContext;
14use crate::protocol::ToolCallResult;
15use crate::render;
16use chrono::Utc;
17use somatize_core::fingerprint::ArchitectureFingerprint;
18use somatize_memory::{ExperimentRecord, RetrievalQuery, derive};
19use somatize_runtime::tracking::{RunReader, summarize};
20use std::path::PathBuf;
21
22/// Retrieve: rank past experiments against a description of the
23/// problem at hand.
24pub fn find_similar(ctx: &SomaContext, params: &serde_json::Value) -> ToolCallResult {
25    let query_text = params
26        .get("query")
27        .and_then(|v| v.as_str())
28        .unwrap_or("")
29        .to_string();
30    let like_run = params.get("like_run").and_then(|v| v.as_str());
31    if query_text.is_empty() && like_run.is_none() {
32        return ToolCallResult::error(
33            "kb_find_similar needs a `query` (free text) or a `like_run` (an experiment id \
34             whose architecture to match), or both.",
35        );
36    }
37
38    let mut query = RetrievalQuery::new(&query_text, Utc::now());
39    query.limit = params
40        .get("limit")
41        .and_then(|v| v.as_u64())
42        .unwrap_or(5)
43        .clamp(1, 50) as usize;
44    if let Some(line) = params.get("research_line").and_then(|v| v.as_str()) {
45        query.research_line = Some(line.to_string());
46    }
47    if let Some(tags) = params.get("tags").and_then(|v| v.as_array()) {
48        query.tags = tags
49            .iter()
50            .filter_map(|t| t.as_str().map(String::from))
51            .collect();
52    }
53    if let Some(days) = params.get("half_life_days").and_then(|v| v.as_f64())
54        && days > 0.0
55    {
56        query.half_life_days = days;
57    }
58
59    if let Some(run_id) = like_run {
60        match architecture_of(ctx, run_id) {
61            Some(architecture) => query.architecture = Some(architecture),
62            None => {
63                return ToolCallResult::error(format!(
64                    "no architecture recorded for '{run_id}' — it may predate fingerprinting, \
65                     or its run directory may be gone. Retry with `query` alone."
66                ));
67            }
68        }
69    }
70
71    match ctx.kb.retrieve(&query) {
72        Ok(hits) => {
73            let label = if query_text.is_empty() {
74                format!("like {}", like_run.unwrap_or_default())
75            } else {
76                query_text
77            };
78            ToolCallResult::text(render::find_similar(&hits, &label))
79        }
80        Err(e) => ToolCallResult::error(format!("retrieval failed: {e}")),
81    }
82}
83
84/// The experiment tree around one run, with the move on every edge.
85pub fn lineage(ctx: &SomaContext, params: &serde_json::Value) -> ToolCallResult {
86    let Some(id) = params.get("id").and_then(|v| v.as_str()) else {
87        return ToolCallResult::error("kb_lineage needs an `id`.");
88    };
89    match ctx.kb.lineage(id) {
90        Ok(Some(lineage)) => ToolCallResult::text(render::lineage(&lineage)),
91        Ok(None) => ToolCallResult::error(unknown_id(id)),
92        Err(e) => ToolCallResult::error(format!("lineage failed: {e}")),
93    }
94}
95
96/// Compare any two experiments — related or not.
97pub fn diff(ctx: &SomaContext, params: &serde_json::Value) -> ToolCallResult {
98    let (Some(a_id), Some(b_id)) = (
99        params.get("a").and_then(|v| v.as_str()),
100        params.get("b").and_then(|v| v.as_str()),
101    ) else {
102        return ToolCallResult::error("kb_diff needs two experiment ids, `a` and `b`.");
103    };
104    let (Some(a), Some(b)) = (fetch(ctx, a_id), fetch(ctx, b_id)) else {
105        let missing = if fetch(ctx, a_id).is_none() {
106            a_id
107        } else {
108            b_id
109        };
110        return ToolCallResult::error(unknown_id(missing));
111    };
112    // The same pure diff the capture path uses, so an on-demand
113    // comparison and a recorded derivation can never disagree.
114    let move_ = derive(&a, &b);
115    ToolCallResult::text(render::diff(&a, &b, &move_))
116}
117
118/// Retain: append a conclusion to an existing experiment.
119///
120/// Written as a separate `Amendment` line. The journal is strictly
121/// append-only: an earlier record is never rewritten, so a note added
122/// today cannot corrupt what was recorded when the run happened.
123pub fn record_conclusion(ctx: &mut SomaContext, params: &serde_json::Value) -> ToolCallResult {
124    let Some(run_id) = params.get("run_id").and_then(|v| v.as_str()) else {
125        return ToolCallResult::error("kb_record_conclusion needs a `run_id`.");
126    };
127    let Some(notes) = params.get("notes").and_then(|v| v.as_str()) else {
128        return ToolCallResult::error(
129            "kb_record_conclusion needs `notes` — what you concluded, in your own words.",
130        );
131    };
132    let Some(target) = fetch(ctx, run_id) else {
133        return ToolCallResult::error(unknown_id(run_id));
134    };
135
136    let id = somatize_core::util::timestamp_id("amend");
137    let mut amendment = ExperimentRecord::amendment(&id, run_id, notes);
138    amendment.research_line = target.research_line.clone();
139    if let Some(hypothesis) = params.get("hypothesis").and_then(|v| v.as_str()) {
140        amendment = amendment.with_hypothesis(hypothesis);
141    }
142    if let Some(tags) = params.get("tags").and_then(|v| v.as_array()) {
143        amendment = amendment.with_tags(
144            tags.iter()
145                .filter_map(|t| t.as_str().map(String::from))
146                .collect(),
147        );
148    }
149
150    match ctx.kb.record(amendment) {
151        Ok(()) => ToolCallResult::text(render::conclusion_recorded(&id, &target)),
152        Err(e) => ToolCallResult::error(format!("failed to record the conclusion: {e}")),
153    }
154}
155
156/// Revise: point `.soma/HEAD` at a run so the next one branches from it.
157pub fn branch_from(ctx: &SomaContext, params: &serde_json::Value) -> ToolCallResult {
158    let Some(run_id) = params.get("run_id").and_then(|v| v.as_str()) else {
159        return ToolCallResult::error("kb_branch_from needs a `run_id`.");
160    };
161    let root = ctx.tracking_root();
162    match somatize_runtime::tracking::checkout(&root, run_id) {
163        Ok(()) => ToolCallResult::text(render::branched(run_id, fetch(ctx, run_id).as_ref())),
164        Err(e) => ToolCallResult::error(format!(
165            "{e}\n\nHEAD was not moved. `kb_stats` lists what this project has recorded."
166        )),
167    }
168}
169
170/// Summarize a run directory on demand.
171///
172/// Reads the directory rather than the journal, so it works on runs
173/// recorded long before the experiment pool existed — and on runs that
174/// never got a journal line at all, because they crashed.
175pub fn summarize_run(ctx: &SomaContext, params: &serde_json::Value) -> ToolCallResult {
176    let Some(run_id) = params.get("run_id").and_then(|v| v.as_str()) else {
177        return ToolCallResult::error("kb_summarize_run needs a `run_id` or a run directory path.");
178    };
179    let Some(dir) = resolve_run_dir(ctx, run_id) else {
180        return ToolCallResult::error(format!(
181            "no run directory for '{run_id}' under {}. Pass an absolute path if the run \
182             lives elsewhere.",
183            ctx.tracking_root().join("runs").display()
184        ));
185    };
186    let summary = RunReader::open(&dir).and_then(|reader| summarize(&reader));
187    match summary {
188        Ok(summary) => ToolCallResult::text(render::summarize_run(&summary)),
189        Err(e) => ToolCallResult::error(format!("cannot read {}: {e}", dir.display())),
190    }
191}
192
193/// Orientation: how big the pool is and how much of it is usable.
194pub fn stats(ctx: &SomaContext, _params: &serde_json::Value) -> ToolCallResult {
195    let records = match ctx.kb.all() {
196        Ok(records) => records,
197        Err(e) => return ToolCallResult::error(format!("cannot read the pool: {e}")),
198    };
199    let lines = ctx.kb.research_lines().unwrap_or_default();
200    ToolCallResult::text(render::stats(
201        &records,
202        &lines,
203        ctx.kb_location().as_deref(),
204    ))
205}
206
207// ── helpers ─────────────────────────────────────────────────────────
208
209fn fetch(ctx: &SomaContext, id: &str) -> Option<ExperimentRecord> {
210    ctx.kb.get(id).ok().flatten()
211}
212
213fn unknown_id(id: &str) -> String {
214    format!(
215        "no experiment '{id}' in this pool. `kb_find_similar` searches by text; `kb_stats` \
216         says how much has been recorded at all."
217    )
218}
219
220/// The architecture to match against: from the record if it has one,
221/// else from the run directory's `fingerprint.json`.
222fn architecture_of(ctx: &SomaContext, run_id: &str) -> Option<ArchitectureFingerprint> {
223    if let Some(architecture) = fetch(ctx, run_id).and_then(|r| r.architecture) {
224        return Some(architecture);
225    }
226    let dir = resolve_run_dir(ctx, run_id)?;
227    let reader = RunReader::open(dir).ok()?;
228    summarize(&reader).ok()?.architecture
229}
230
231/// Accept a run id, a path relative to the project, or an absolute path.
232fn resolve_run_dir(ctx: &SomaContext, run_id: &str) -> Option<PathBuf> {
233    let as_path = PathBuf::from(run_id);
234    if as_path.is_dir() {
235        return Some(as_path);
236    }
237    if let Some(dir) = fetch(ctx, run_id).and_then(|r| r.run_dir) {
238        let dir = PathBuf::from(dir);
239        if dir.is_dir() {
240            return Some(dir);
241        }
242    }
243    let under_root = ctx.tracking_root().join("runs").join(run_id);
244    under_root.is_dir().then_some(under_root)
245}