somatize_mcp/tools/
mod.rs

1//! Tool definitions and dispatch for the Soma MCP server.
2//!
3//! Descriptions here are the only documentation a model gets, so they
4//! say what a tool actually does — including when that is "not much".
5
6pub mod knowledge;
7
8use crate::context::SomaContext;
9use crate::protocol::{ToolCallResult, ToolDefinition};
10use serde_json::json;
11
12/// Register all available tools.
13pub fn all_tools() -> Vec<ToolDefinition> {
14    vec![
15        // ── Code tools ──
16        ToolDefinition {
17            name: "list_filters".into(),
18            description: "List available filter source files in the project directory.".into(),
19            input_schema: json!({
20                "type": "object",
21                "properties": {
22                    "path": { "type": "string", "description": "Project directory path (optional, uses configured default)" }
23                }
24            }),
25        },
26        ToolDefinition {
27            name: "read_filter_source".into(),
28            description: "Read the source code of a filter file.".into(),
29            input_schema: json!({
30                "type": "object",
31                "properties": {
32                    "file_path": { "type": "string", "description": "Path to the filter source file" }
33                },
34                "required": ["file_path"]
35            }),
36        },
37        ToolDefinition {
38            name: "write_filter_source".into(),
39            description: "Write or update filter source code. Creates a backup before writing."
40                .into(),
41            input_schema: json!({
42                "type": "object",
43                "properties": {
44                    "file_path": { "type": "string", "description": "Path to the filter source file" },
45                    "content": { "type": "string", "description": "New file content" }
46                },
47                "required": ["file_path", "content"]
48            }),
49        },
50        // ── Execution tools ──
51        ToolDefinition {
52            name: "run_pipeline".into(),
53            description: "Build a computation graph out of the project's filters and RUN it. \
54                 Each node names a filter as `module.Class`, `path/to/file.py:Class`, or a \
55                 bare class name found in the files list_filters returns; `config` is its \
56                 constructor keywords. Edges connect node ids. The graph is fitted and then \
57                 forwarded on `input`, in a Python subprocess rooted at the project, so the \
58                 filters you just read with read_filter_source are the ones that run. \
59                 Tracked by default: the result carries a run_dir, and kb_summarize_run will \
60                 read it back. This EXECUTES project code."
61                .into(),
62            input_schema: json!({
63                "type": "object",
64                "properties": {
65                    "nodes": {
66                        "type": "array",
67                        "description": "Graph nodes, in any order",
68                        "items": {
69                            "type": "object",
70                            "properties": {
71                                "id": { "type": "string", "description": "Node id, used by edges" },
72                                "filter": { "type": "string", "description": "module.Class | path.py:Class | ClassName" },
73                                "config": { "type": "object", "description": "Constructor keyword arguments" },
74                                "target": { "type": "string", "description": "Worker tag, for a distributed node" }
75                            },
76                            "required": ["id", "filter"]
77                        }
78                    },
79                    "edges": {
80                        "type": "array",
81                        "description": "Directed edges as [from_id, to_id] pairs",
82                        "items": { "type": "array", "items": { "type": "string" } }
83                    },
84                    "input": { "description": "Input data: a number, a list, a nested list, or an object" },
85                    "y": { "description": "Targets, for a supervised fit" },
86                    "fit": { "type": "boolean", "description": "Fit before forwarding (default true)" },
87                    "cache": { "type": "string", "enum": ["memory", "tiered", "none"], "description": "Cache backend (default memory)" },
88                    "track": { "type": "boolean", "description": "Record into the experiment pool (default true)" },
89                    "name": { "type": "string", "description": "Run name, as it appears in the pool" },
90                    "tags": { "type": "array", "items": { "type": "string" } },
91                    "params": { "type": "object", "description": "Hyperparameters that live outside the graph, recorded with the run" }
92                },
93                "required": ["nodes", "input"]
94            }),
95        },
96        ToolDefinition {
97            name: "run_study".into(),
98            description: "Search a graph's hyperparameters and RUN the trials. Same node spec \
99                 as run_pipeline, with one difference: any config value written as \
100                 {\"__search__\": {\"low\": 1e-4, \"high\": 1e-1, \"scale\": \"log\"}} or \
101                 {\"__search__\": {\"choices\": [...]}} becomes a dimension to search. The \
102                 graph is rebuilt per trial with the sampled values, fitted, and forwarded; \
103                 its output is the objective — a number, or an object from which `metric` is \
104                 read (soma.library.Eval emits one). Returns the best trial and a run_dir. \
105                 This EXECUTES project code, n_trials times."
106                .into(),
107            input_schema: json!({
108                "type": "object",
109                "properties": {
110                    "nodes": {
111                        "type": "array",
112                        "description": "Graph nodes; mark searched values with {\"__search__\": {...}}",
113                        "items": {
114                            "type": "object",
115                            "properties": {
116                                "id": { "type": "string" },
117                                "filter": { "type": "string" },
118                                "config": { "type": "object" }
119                            },
120                            "required": ["id", "filter"]
121                        }
122                    },
123                    "edges": { "type": "array", "items": { "type": "array", "items": { "type": "string" } } },
124                    "input": { "description": "Input data, the same for every trial" },
125                    "y": { "description": "Targets, for a supervised fit" },
126                    "name": { "type": "string", "description": "Study name" },
127                    "strategy": { "type": "string", "enum": ["grid", "random", "bayesian"], "description": "Sampler (default random)" },
128                    "n_trials": { "type": "integer", "description": "How many trials to run (default 10)" },
129                    "metric": { "type": "string", "description": "Key to read from the graph's output when it is an object (default \"score\")" },
130                    "direction": { "type": "string", "enum": ["minimize", "maximize"], "description": "Default minimize" },
131                    "seed": { "type": "integer", "description": "Seeds the sampler, so the search repeats" },
132                    "cache": { "type": "string", "enum": ["memory", "tiered", "none"] }
133                },
134                "required": ["nodes", "input"]
135            }),
136        },
137        // ── Knowledge tools ──
138        ToolDefinition {
139            name: "record_experiment".into(),
140            description: "Record an experiment by hand. Runs started with graph.track_run() or \
141                 study.run() record themselves — with a conclusion, an architecture fingerprint \
142                 and a lineage — so use this only for work soma did not execute. To add a \
143                 finding to an existing experiment, use kb_record_conclusion instead."
144                .into(),
145            input_schema: json!({
146                "type": "object",
147                "properties": {
148                    "id": { "type": "string" },
149                    "name": { "type": "string" },
150                    "hypothesis": { "type": "string", "description": "What you expected, and why" },
151                    "research_line": { "type": "string", "description": "Groups related experiments; inherited from the parent when one is given" },
152                    "pipeline_summary": { "type": "string", "description": "One line describing the topology, e.g. 'scaler → encoder → head'" },
153                    "params": { "type": "object", "description": "Hyperparameters; a later variant diffs against these" },
154                    "metrics": { "type": "object", "description": "Final numeric results" },
155                    "parent": { "type": "string", "description": "Experiment this one was derived from. Setting it computes the move between them." },
156                    "objective": { "type": "string", "description": "Metric being optimized, so improvement can be judged" },
157                    "run_dir": { "type": "string", "description": "Directory holding raw artifacts, if any" },
158                    "tags": { "type": "array", "items": { "type": "string" } },
159                    "notes": { "type": "string", "description": "What you concluded" }
160                },
161                "required": ["id", "name"]
162            }),
163        },
164        // ── Experiment pool ──
165        ToolDefinition {
166            name: "kb_find_similar".into(),
167            description: "Find past experiments bearing on the problem at hand — the first \
168                 thing to call before designing anything. Ranks by text relevance, \
169                 architectural resemblance, recency and importance. Dead ends rank too: not \
170                 repeating a failure saves as much time as repeating a success. Each hit \
171                 carries its conclusion, the move that produced it, and a run_dir you can read \
172                 directly."
173                .into(),
174            input_schema: json!({
175                "type": "object",
176                "properties": {
177                    "query": { "type": "string", "description": "What you are trying to do, in words. Names, filters, metrics and symptoms all match." },
178                    "like_run": { "type": "string", "description": "Experiment id whose architecture to match. Combine with query, or use alone to find structurally similar work." },
179                    "limit": { "type": "integer", "default": 5, "description": "1-50" },
180                    "research_line": { "type": "string", "description": "Restrict to one line" },
181                    "tags": { "type": "array", "items": { "type": "string" }, "description": "Every listed tag must be present" },
182                    "half_life_days": { "type": "number", "default": 30, "description": "Raise it to weight old work more heavily" }
183                }
184            }),
185        },
186        ToolDefinition {
187            name: "kb_lineage".into(),
188            description: "The experiment tree around one run: ancestors above, descendants \
189                 below, and on every edge the change that produced the child from its parent \
190                 plus what it did to the metrics. This is how you see what has already been \
191                 tried from a given starting point."
192                .into(),
193            input_schema: json!({
194                "type": "object",
195                "properties": {
196                    "id": { "type": "string", "description": "Experiment or run id" }
197                },
198                "required": ["id"]
199            }),
200        },
201        ToolDefinition {
202            name: "kb_diff".into(),
203            description: "Compare two experiments: what differs in architecture, parameters \
204                 and code, what each metric did, and what it cost (wall time, cache hits). \
205                 They need not be related — this works on any two ids."
206                .into(),
207            input_schema: json!({
208                "type": "object",
209                "properties": {
210                    "a": { "type": "string", "description": "Baseline experiment id" },
211                    "b": { "type": "string", "description": "Variant experiment id" }
212                },
213                "required": ["a", "b"]
214            }),
215        },
216        ToolDefinition {
217            name: "kb_record_conclusion".into(),
218            description: "Retain what you learned about a run: why it worked, why it did not, \
219                 what to try next. Appended as a separate amendment — the original record is \
220                 never rewritten — and indexed so a later kb_find_similar surfaces it. Worth \
221                 doing for failures especially."
222                .into(),
223            input_schema: json!({
224                "type": "object",
225                "properties": {
226                    "run_id": { "type": "string", "description": "Experiment being annotated" },
227                    "notes": { "type": "string", "description": "What you concluded" },
228                    "hypothesis": { "type": "string", "description": "The hypothesis this run turned out to be testing" },
229                    "tags": { "type": "array", "items": { "type": "string" } }
230                },
231                "required": ["run_id", "notes"]
232            }),
233        },
234        ToolDefinition {
235            name: "kb_branch_from".into(),
236            description: "Point .soma/HEAD at an existing run so the NEXT run records itself \
237                 as its child. Use it to go back and try a different variation instead of \
238                 continuing from the last thing that happened to run. Creates a sibling \
239                 branch; it never moves or rewrites existing history."
240                .into(),
241            input_schema: json!({
242                "type": "object",
243                "properties": {
244                    "run_id": { "type": "string", "description": "Run to branch from. Must exist under .soma/runs/." }
245                },
246                "required": ["run_id"]
247            }),
248        },
249        ToolDefinition {
250            name: "kb_summarize_run".into(),
251            description: "Read a run directory and summarize it: outcome, metrics, slowest \
252                 node, cache effectiveness, health flags, trials. Works on runs recorded \
253                 before the experiment pool existed and on runs that crashed before writing a \
254                 journal line, and reports what it could not read rather than staying silent."
255                .into(),
256            input_schema: json!({
257                "type": "object",
258                "properties": {
259                    "run_id": { "type": "string", "description": "Run id under .soma/runs/, or a path to a run directory" }
260                },
261                "required": ["run_id"]
262            }),
263        },
264        ToolDefinition {
265            name: "kb_stats".into(),
266            description: "How much this project has recorded and how usable it is: totals, \
267                 date span, research lines, and honest coverage — what fraction of records \
268                 carry a conclusion, a lineage, an architecture. Call it first when you do not \
269                 know whether the pool is worth querying."
270                .into(),
271            input_schema: json!({ "type": "object", "properties": {} }),
272        },
273        ToolDefinition {
274            name: "query_knowledge_base".into(),
275            description: "Search experiments in the knowledge base by text query.".into(),
276            input_schema: json!({
277                "type": "object",
278                "properties": {
279                    "query": { "type": "string", "description": "Search query" },
280                    "max_results": { "type": "integer", "default": 10 }
281                },
282                "required": ["query"]
283            }),
284        },
285        ToolDefinition {
286            name: "get_trajectory".into(),
287            description: "Get the metric trajectory for a research line.".into(),
288            input_schema: json!({
289                "type": "object",
290                "properties": {
291                    "research_line": { "type": "string" },
292                    "metric": { "type": "string" }
293                },
294                "required": ["research_line", "metric"]
295            }),
296        },
297        ToolDefinition {
298            name: "get_change_points".into(),
299            description: "Detect significant changes in experiment metrics.".into(),
300            input_schema: json!({
301                "type": "object",
302                "properties": {
303                    "research_line": { "type": "string" },
304                    "metric": { "type": "string" },
305                    "threshold": { "type": "number", "default": 0.05 }
306                },
307                "required": ["research_line", "metric"]
308            }),
309        },
310        ToolDefinition {
311            name: "list_research_lines".into(),
312            description: "List all research lines with trend analysis.".into(),
313            input_schema: json!({
314                "type": "object",
315                "properties": {}
316            }),
317        },
318        ToolDefinition {
319            name: "promising_lines".into(),
320            description: "Get research lines that are improving.".into(),
321            input_schema: json!({
322                "type": "object",
323                "properties": {
324                    "metric": { "type": "string", "description": "Metric to evaluate" }
325                },
326                "required": ["metric"]
327            }),
328        },
329        // ── Project tools ──
330        ToolDefinition {
331            name: "create_research_line".into(),
332            description: "Create a new research line for tracking experiments.".into(),
333            input_schema: json!({
334                "type": "object",
335                "properties": {
336                    "name": { "type": "string" },
337                    "description": { "type": "string" }
338                },
339                "required": ["name"]
340            }),
341        },
342        ToolDefinition {
343            name: "generate_report".into(),
344            description: "Generate a markdown report for a research line.".into(),
345            input_schema: json!({
346                "type": "object",
347                "properties": {
348                    "research_line": { "type": "string" }
349                },
350                "required": ["research_line"]
351            }),
352        },
353    ]
354}
355
356/// Dispatch a tool call to the appropriate handler.
357///
358/// Every knowledge read is preceded by a refresh: this server outlives
359/// the training runs it is asked about, and a stale snapshot silently
360/// answers "no such experiment" for a run that finished five minutes
361/// ago in another terminal.
362pub fn dispatch(
363    ctx: &mut SomaContext,
364    tool_name: &str,
365    params: &serde_json::Value,
366) -> ToolCallResult {
367    if reads_knowledge(tool_name) {
368        ctx.refresh_kb();
369    }
370    match tool_name {
371        // Experiment pool
372        "kb_find_similar" => knowledge::find_similar(ctx, params),
373        "kb_lineage" => knowledge::lineage(ctx, params),
374        "kb_diff" => knowledge::diff(ctx, params),
375        "kb_record_conclusion" => knowledge::record_conclusion(ctx, params),
376        "kb_branch_from" => knowledge::branch_from(ctx, params),
377        "kb_summarize_run" => knowledge::summarize_run(ctx, params),
378        "kb_stats" => knowledge::stats(ctx, params),
379        // Code tools
380        "list_filters" => ctx.list_filters(params),
381        "read_filter_source" => ctx.read_filter_source(params),
382        "write_filter_source" => ctx.write_filter_source(params),
383        // Execution tools
384        "run_pipeline" => ctx.run_pipeline(params),
385        "run_study" => ctx.run_study(params),
386        // Knowledge tools
387        "record_experiment" => ctx.record_experiment(params),
388        "query_knowledge_base" => ctx.query_knowledge_base(params),
389        "get_trajectory" => ctx.get_trajectory(params),
390        "get_change_points" => ctx.get_change_points(params),
391        "list_research_lines" => ctx.list_research_lines(params),
392        "promising_lines" => ctx.promising_lines(params),
393        // Project tools
394        "create_research_line" => ctx.create_research_line(params),
395        "generate_report" => ctx.generate_report(params),
396        _ => ToolCallResult::error(format!("Unknown tool: {tool_name}")),
397    }
398}
399
400/// Whether a tool reads the experiment journal, and therefore needs to
401/// see what other processes have appended.
402fn reads_knowledge(tool_name: &str) -> bool {
403    tool_name.starts_with("kb_")
404        || matches!(
405            tool_name,
406            "query_knowledge_base"
407                | "get_trajectory"
408                | "get_change_points"
409                | "list_research_lines"
410                | "promising_lines"
411                | "generate_report"
412                | "record_experiment"
413        )
414}
415
416#[cfg(test)]
417mod tests {
418    use super::*;
419
420    #[test]
421    fn every_tool_is_defined_exactly_once_and_dispatches() {
422        let tools = all_tools();
423        let mut names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
424        names.sort();
425        let unique = {
426            let mut u = names.clone();
427            u.dedup();
428            u
429        };
430        assert_eq!(names, unique, "duplicate tool name");
431
432        let dir = tempfile::tempdir().unwrap();
433        for tool in &tools {
434            let mut ctx = SomaContext::with_env_override(dir.path(), None);
435            let result = dispatch(&mut ctx, &tool.name, &serde_json::json!({}));
436            // No arguments: a tool may refuse, but never as "unknown".
437            let text = result.content_text();
438            assert!(
439                !text.contains("Unknown tool"),
440                "{} is defined but not dispatched",
441                tool.name
442            );
443        }
444        assert!(
445            dispatch(
446                &mut SomaContext::with_env_override(dir.path(), None),
447                "kb_nope",
448                &serde_json::json!({})
449            )
450            .content_text()
451            .contains("Unknown tool")
452        );
453    }
454
455    #[test]
456    fn the_pool_tools_are_all_registered() {
457        let names: Vec<String> = all_tools().into_iter().map(|t| t.name).collect();
458        for expected in [
459            "kb_find_similar",
460            "kb_lineage",
461            "kb_diff",
462            "kb_record_conclusion",
463            "kb_branch_from",
464            "kb_summarize_run",
465            "kb_stats",
466        ] {
467            assert!(names.contains(&expected.to_string()), "missing {expected}");
468        }
469        assert_eq!(names.len(), 20, "13 original tools + 7 pool tools");
470    }
471
472    #[test]
473    fn the_execution_tools_say_that_they_execute() {
474        // They spent a long time declaring "NOT IMPLEMENTED" and echoing
475        // their arguments. They run now — and a tool that runs project
476        // code has to say so where the model reads it, not only in a
477        // design document.
478        for name in ["run_pipeline", "run_study"] {
479            let tool = all_tools().into_iter().find(|t| t.name == name).unwrap();
480            assert!(
481                !tool.description.contains("NOT IMPLEMENTED"),
482                "{name} still claims to do nothing: {}",
483                tool.description
484            );
485            assert!(
486                tool.description.contains("EXECUTES project code"),
487                "{name} must announce that it executes: {}",
488                tool.description
489            );
490            // `nodes` is what turns these from a vague "config" into a
491            // graph a model can actually describe.
492            let props = tool.input_schema.get("properties").unwrap();
493            assert!(props.get("nodes").is_some(), "{name} takes no nodes");
494            assert!(props.get("input").is_some(), "{name} takes no input");
495        }
496    }
497
498    #[test]
499    fn knowledge_reads_are_refreshed_and_code_tools_are_not() {
500        for name in ["kb_find_similar", "query_knowledge_base", "generate_report"] {
501            assert!(reads_knowledge(name), "{name} should refresh");
502        }
503        for name in ["list_filters", "read_filter_source", "run_pipeline"] {
504            assert!(!reads_knowledge(name), "{name} should not refresh");
505        }
506    }
507}
somatize_mcp/tools/mod.rs

somatize_mcp/tools/
mod.rs
