somatize_mcp/tools/
knowledge.rs1use 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
22pub 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
84pub 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
96pub 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 let move_ = derive(&a, &b);
115 ToolCallResult::text(render::diff(&a, &b, &move_))
116}
117
118pub 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
156pub 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
170pub 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
193pub 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
207fn 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
220fn 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
231fn 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}