1use crate::index::SearchResult;
2use noether_core::stage::Stage;
3use noether_core::types::NType;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct SynthesisSpec {
12 pub name: String,
13 pub description: String,
14 pub input: NType,
15 pub output: NType,
16 pub rationale: String,
17}
18
19#[derive(Debug, Clone, Deserialize)]
21pub struct SynthesisResponse {
22 pub examples: Vec<SynthesisExample>,
23 pub implementation: String,
24 #[serde(default = "default_language")]
25 pub language: String,
26}
27
28#[derive(Debug, Clone, Deserialize)]
29pub struct SynthesisExample {
30 pub input: Value,
31 pub output: Value,
32}
33
34fn default_language() -> String {
35 "python".into()
36}
37
38pub fn build_system_prompt(candidates: &[(&SearchResult, &Stage)]) -> String {
41 let mut prompt = String::new();
42
43 prompt.push_str(
45 "You are Noether's Composition Agent. You translate problem descriptions into \
46 composition graphs in Lagrange JSON format.\n\n",
47 );
48
49 prompt.push_str("## CRITICAL RULES\n\n");
51 prompt.push_str("1. ONLY use stage IDs from the AVAILABLE STAGES list. Never invent IDs.\n");
52 prompt.push_str("2. Types MUST match: the output type of one stage must be a subtype of the next stage's input type.\n");
53 prompt.push_str("3. Most stages take Record inputs with SPECIFIC FIELD NAMES. If a stage needs Record{items,key,...} but your pipeline produces a bare List, DO NOT try Parallel+Const wiring — SYNTHESIZE a stage instead.\n");
54 prompt.push_str("4. Output ONLY a JSON code block — no explanation before or after.\n");
55 prompt.push_str("5. EVERY node in the graph (including nested ones) MUST have an `\"op\"` field. There are NO exceptions.\n");
56 prompt.push_str(" Valid values: `\"Stage\"`, `\"Const\"`, `\"Sequential\"`, `\"Parallel\"`, `\"Branch\"`, `\"Fanout\"`, `\"Retry\"`.\n");
57 prompt.push_str("6. NEVER use a Stage branch in Parallel to \"pass through\" the input. Parallel branches receive the input but Stage branches transform it. Use Const for literal values only.\n\n");
58
59 prompt.push_str("## Type System\n\n");
61 prompt
62 .push_str("- `Any` accepts any value. `Text`, `Number`, `Bool`, `Null` are primitives.\n");
63 prompt.push_str("- `Record { field: Type }` is an object with named fields. The stage REQUIRES exactly those fields.\n");
64 prompt.push_str("- `List<T>` is an array. `Map<K,V>` is a key-value object.\n");
65 prompt.push_str("- `T | Null` means the field is optional (can be null).\n");
66 prompt.push_str(
67 "- Width subtyping: `{a, b, c}` is subtype of `{a, b}` — extra fields are OK.\n\n",
68 );
69
70 prompt.push_str("## Operators\n\n");
72 prompt.push_str("- **Stage**: `{\"op\": \"Stage\", \"id\": \"<hash>\"}` — optionally add `\"config\": {\"key\": \"value\"}` to provide static parameters\n");
73 prompt.push_str("- **Const**: `{\"op\": \"Const\", \"value\": <any JSON value>}` — emits a literal constant, ignores its input entirely\n");
74 prompt.push_str("- **Sequential**: `{\"op\": \"Sequential\", \"stages\": [A, B, C]}` — output of A feeds B, then C\n");
75 prompt.push_str("- **Parallel**: `{\"op\": \"Parallel\", \"branches\": {\"key1\": A, \"key2\": B}}` — ALL branches receive the SAME full input (or the field matching the branch name if the input is a Record); outputs are merged into a Record `{\"key1\": <out_A>, \"key2\": <out_B>}`\n");
76 prompt.push_str("- **Branch**: `{\"op\": \"Branch\", \"predicate\": P, \"if_true\": A, \"if_false\": B}` — P receives the original input and MUST return Bool; A and B also receive the SAME original input (NOT the Bool)\n");
77 prompt.push_str("- **Fanout**: `{\"op\": \"Fanout\", \"source\": A, \"targets\": [B, C]}`\n");
78 prompt.push_str("- **Retry**: `{\"op\": \"Retry\", \"stage\": A, \"max_attempts\": 3, \"delay_ms\": 500}`\n\n");
79
80 prompt.push_str("## Stage Config — VERY IMPORTANT\n\n");
82 prompt.push_str(
83 "Many stages need `Record { items: List, key: Text, descending: Bool }` as input.\n",
84 );
85 prompt.push_str("The pipeline only provides the DATA (e.g., a `List`). The PARAMETERS (`key`, `descending`) are static.\n\n");
86 prompt.push_str("**Use `config` to supply static parameters.** The executor merges config fields with the pipeline input:\n\n");
87 prompt.push_str("```json\n");
88 prompt.push_str("{\n");
89 prompt.push_str(" \"op\": \"Stage\",\n");
90 prompt.push_str(" \"id\": \"<list_sort_id>\",\n");
91 prompt.push_str(" \"config\": {\"key\": \"score\", \"descending\": true}\n");
92 prompt.push_str("}\n");
93 prompt.push_str("```\n\n");
94 prompt.push_str("The pipeline flows `List<Any>` → the executor produces `{items: <the list>, key: \"score\", descending: true}` → `list_sort` receives exactly what it needs.\n\n");
95 prompt.push_str("**Rules for config:**\n");
96 prompt.push_str("- Use config for PARAMETER fields (key, count, delimiter, pattern, etc.)\n");
97 prompt.push_str("- The pipeline provides the DATA field (items, text, data, records, etc.)\n");
98 prompt.push_str("- Config keys must match the stage's Record field names exactly\n");
99 prompt.push_str("- Config values are JSON literals (strings, numbers, booleans, null)\n\n");
100 prompt.push_str("**Example: CSV parse → sort by revenue → take top 3 → serialize**\n");
101 prompt.push_str("```json\n");
102 prompt.push_str("{\n");
103 prompt.push_str(" \"op\": \"Sequential\",\n");
104 prompt.push_str(" \"stages\": [\n");
105 prompt.push_str(" {\"op\": \"Stage\", \"id\": \"<csv_parse_id>\"},\n");
106 prompt.push_str(" {\"op\": \"Stage\", \"id\": \"<list_sort_id>\", \"config\": {\"key\": \"revenue\", \"descending\": true}},\n");
107 prompt.push_str(
108 " {\"op\": \"Stage\", \"id\": \"<list_take_id>\", \"config\": {\"count\": 3}},\n",
109 );
110 prompt.push_str(" {\"op\": \"Stage\", \"id\": \"<json_serialize_id>\"}\n");
111 prompt.push_str(" ]\n");
112 prompt.push_str("}\n");
113 prompt.push_str("```\n\n");
114 prompt.push_str("**Parallel** is still used for running branches concurrently on the same input — NOT for assembling Record parameters.\n\n");
115
116 prompt.push_str("## Branch Operator — How It Works\n\n");
118 prompt.push_str("```\nBranch receives input X.\n");
119 prompt.push_str("1. Runs predicate(X) → must return Bool\n");
120 prompt.push_str("2. If true: runs if_true(X) — same X, NOT the Bool\n");
121 prompt.push_str("3. If false: runs if_false(X) — same X, NOT the Bool\n```\n\n");
122 prompt.push_str("Do NOT use Branch when you mean a stage that selects between values.\n");
123 prompt.push_str(
124 "Branch is for routing execution to different sub-graphs based on a condition.\n\n",
125 );
126
127 prompt.push_str("## When to Synthesize a New Stage\n\n");
129 prompt.push_str("**PREFER SYNTHESIS over complex composition** in these cases:\n\n");
130 prompt.push_str("- The required primitive operation (e.g. modulo, even/odd, filter, sort-by-key) has no matching stage.\n");
131 prompt.push_str("- Solving the problem would need 3+ stages of awkward Record manipulation.\n");
132 prompt.push_str("- You need to filter a list, transform each element with custom logic, or reshape data in a bespoke way.\n");
133 prompt.push_str(
134 "- **You need to call a SPECIFIC external HTTP API** — always synthesize for API calls.\n",
135 );
136 prompt.push_str(" The `http_get` stdlib stage is for generic URL fetching; it cannot parse JSON, extract fields,\n");
137 prompt.push_str(" or format results specific to a given API. Always synthesize a stage that does the full\n");
138 prompt.push_str(" HTTP call + parse + reshape in one Python function.\n\n");
139 prompt.push_str("**CRITICAL: A synthesis request is a STANDALONE top-level document.**\n");
140 prompt.push_str(
141 "It CANNOT be embedded inside a `Sequential.stages` list or any other graph node.\n",
142 );
143 prompt.push_str("You MUST choose ONE of these two responses per turn:\n");
144 prompt.push_str(" Option A) A synthesis request (to register a missing stage), OR\n");
145 prompt.push_str(
146 " Option B) A composition graph (using existing + already-registered stages).\n",
147 );
148 prompt.push_str(
149 "If you return a synthesis request, the stage will be registered and you WILL get\n",
150 );
151 prompt
152 .push_str("another turn to compose using that stage. Do NOT mix them in one response.\n\n");
153 prompt
154 .push_str("**Synthesis format (respond with ONLY this — no graph, no explanation):**\n\n");
155 prompt.push_str("```json\n");
156 prompt.push_str("{\n");
157 prompt.push_str(" \"action\": \"synthesize\",\n");
158 prompt.push_str(" \"spec\": {\n");
159 prompt.push_str(" \"name\": \"snake_case_stage_name\",\n");
160 prompt.push_str(" \"description\": \"One-sentence description of what this stage does\",\n");
161 prompt.push_str(" \"input\": {\"kind\": \"Text\"},\n");
162 prompt.push_str(" \"output\": {\"kind\": \"Number\"},\n");
163 prompt.push_str(" \"rationale\": \"Why no available stage satisfies this\"\n");
164 prompt.push_str(" }\n");
165 prompt.push_str("}\n");
166 prompt.push_str("```\n\n");
167 prompt.push_str("NType JSON format:\n");
168 prompt.push_str("- Primitives: `{\"kind\":\"Text\"}`, `{\"kind\":\"Number\"}`, `{\"kind\":\"Bool\"}`, `{\"kind\":\"Any\"}`, `{\"kind\":\"Null\"}`\n");
169 prompt.push_str("- List: `{\"kind\":\"List\",\"value\":<T>}`\n");
170 prompt.push_str("- Map: `{\"kind\":\"Map\",\"value\":{\"key\":{\"kind\":\"Text\"},\"value\":<T>}}` ← note: Map.value is an object with `key` and `value` fields\n");
171 prompt.push_str("- Record: `{\"kind\":\"Record\",\"value\":{\"field_name\":<T>,...}}`\n");
172 prompt.push_str("- Union: `{\"kind\":\"Union\",\"value\":[<T>,...]}`\n\n");
173 prompt.push_str("**Keep synthesis types SIMPLE:**\n");
174 prompt.push_str(
175 "- Use `Any` for complex or heterogeneous output (lists of dicts, nested structures).\n",
176 );
177 prompt.push_str("- Use `Text` for input when it's raw data (CSV text, JSON string).\n");
178 prompt.push_str("- Do NOT use `Map<Text, Any>` — use `Any` instead.\n");
179 prompt.push_str(
180 "- Prefer flat types: `Text → Any`, `Record{text: Text} → Any`, `Any → Text`.\n\n",
181 );
182 prompt.push_str("**Examples that SHOULD use synthesis:**\n");
183 prompt.push_str(
184 "- \"check if a number is even or odd\" → synthesize `is_even_or_odd` (Number → Text)\n",
185 );
186 prompt.push_str("- \"filter a list keeping items that match a pattern\" → synthesize `filter_by_pattern` (Record { items, pattern } → List)\n");
187 prompt.push_str("- \"sort a list by a field\" → synthesize `sort_by_field` (Record { items, field } → List)\n");
188 prompt.push_str("- \"sort a list and take the top N\" → synthesize `sort_and_take` (Record { items, n } → List)\n");
189 prompt.push_str("- \"search npm packages and return results\" → synthesize `npm_search` (Record { query, limit } → List) — NEVER try to compose with http_get\n");
190 prompt.push_str("- \"search GitHub repos\" → synthesize `github_search` — NEVER try to compose with http_get\n");
191 prompt.push_str("- ANY call to a named external API (GitHub, npm, Hacker News, Spotify, etc.) → synthesize\n\n");
192
193 let parse_json_id = find_candidate_id(candidates, "Parse a JSON string");
195 let to_json_id = find_candidate_id(candidates, "Serialize any value to a JSON");
196 let is_null_id = find_candidate_id(candidates, "Check if a value is null");
197 let text_upper_id = find_candidate_id(candidates, "Convert text to uppercase");
198 let text_lower_id = find_candidate_id(candidates, "Convert text to lowercase");
199
200 prompt.push_str("## EXAMPLE 1: Sequential composition\n\n");
201 prompt.push_str("Problem: \"Parse a JSON string and serialize it back\"\n\n");
202 prompt.push_str("The stage `parse_json` has input `Text` and output `Any`.\n");
203 prompt.push_str("The stage `to_json` has input `Any` and output `Text`.\n");
204 prompt.push_str("Since `Any` (output of parse_json) is subtype of `Any` (input of to_json), they compose.\n\n");
205 prompt.push_str("```json\n");
206 prompt.push_str("{\n");
207 prompt.push_str(" \"description\": \"Parse JSON then serialize back to text\",\n");
208 prompt.push_str(" \"version\": \"0.1.0\",\n");
209 prompt.push_str(" \"root\": {\n");
210 prompt.push_str(" \"op\": \"Sequential\",\n");
211 prompt.push_str(" \"stages\": [\n");
212 prompt.push_str(&format!(
213 " {{\"op\": \"Stage\", \"id\": \"{}\"}},\n",
214 parse_json_id
215 ));
216 prompt.push_str(&format!(
217 " {{\"op\": \"Stage\", \"id\": \"{}\"}}\n",
218 to_json_id
219 ));
220 prompt.push_str(" ]\n");
221 prompt.push_str(" }\n");
222 prompt.push_str("}\n");
223 prompt.push_str("```\n\n");
224
225 prompt.push_str("## EXAMPLE 2: Branch operator (condition-based routing)\n\n");
226 prompt.push_str("Problem: \"Convert text to uppercase if it is not null, otherwise return empty string\"\n\n");
227 prompt.push_str(
228 "The `Branch` predicate receives the original `Text | Null` input and returns `Bool`.\n",
229 );
230 prompt.push_str("`if_true` and `if_false` ALSO receive the original input — NOT the Bool.\n\n");
231 prompt.push_str("```json\n");
232 prompt.push_str("{\n");
233 prompt.push_str(" \"description\": \"Uppercase non-null text\",\n");
234 prompt.push_str(" \"version\": \"0.1.0\",\n");
235 prompt.push_str(" \"root\": {\n");
236 prompt.push_str(" \"op\": \"Branch\",\n");
237 prompt.push_str(&format!(
238 " \"predicate\": {{\"op\": \"Stage\", \"id\": \"{}\"}},\n",
239 is_null_id
240 ));
241 prompt.push_str(&format!(
242 " \"if_true\": {{\"op\": \"Stage\", \"id\": \"{}\"}},\n",
243 text_lower_id
244 ));
245 prompt.push_str(&format!(
246 " \"if_false\": {{\"op\": \"Stage\", \"id\": \"{}\"}}\n",
247 text_upper_id
248 ));
249 prompt.push_str(" }\n");
250 prompt.push_str("}\n");
251 prompt.push_str("```\n\n");
252
253 prompt.push_str("## EXAMPLE 3: Const + Parallel to assemble a multi-field Record\n\n");
254 prompt.push_str(
255 "Problem: \"Search for repos, then format a report with a fixed topic and summary\"\n\n",
256 );
257 prompt.push_str("The search stage returns a List. The format stage needs `Record{topic, results, summary}`.\n");
258 prompt.push_str("Use Parallel: `results` branch runs the search (receives full input), `topic` and `summary` are Const literals.\n\n");
259 prompt.push_str("```json\n");
260 prompt.push_str("{\n");
261 prompt.push_str(" \"description\": \"Search then format a report\",\n");
262 prompt.push_str(" \"version\": \"0.1.0\",\n");
263 prompt.push_str(" \"root\": {\n");
264 prompt.push_str(" \"op\": \"Sequential\",\n");
265 prompt.push_str(" \"stages\": [\n");
266 prompt.push_str(" {\n");
267 prompt.push_str(" \"op\": \"Parallel\",\n");
268 prompt.push_str(" \"branches\": {\n");
269 prompt.push_str(" \"results\": {\"op\": \"Stage\", \"id\": \"<search_stage_id>\"},\n");
270 prompt.push_str(" \"topic\": {\"op\": \"Const\", \"value\": \"async runtime\"},\n");
271 prompt.push_str(
272 " \"summary\": {\"op\": \"Const\", \"value\": \"Top async runtime libraries\"}\n",
273 );
274 prompt.push_str(" }\n");
275 prompt.push_str(" },\n");
276 prompt.push_str(" {\"op\": \"Stage\", \"id\": \"<format_stage_id>\"}\n");
277 prompt.push_str(" ]\n");
278 prompt.push_str(" }\n");
279 prompt.push_str("}\n");
280 prompt.push_str("```\n\n");
281
282 let sort_id = find_candidate_id(candidates, "Sort a list");
284 let take_id = find_candidate_id(candidates, "Take the first N elements");
285 let json_ser_id = find_candidate_id(candidates, "Serialize any value to a JSON");
286
287 prompt.push_str("## EXAMPLE 4: Using config for parameterized stages\n\n");
288 prompt.push_str("Problem: \"Sort a list by score descending and take the top 3\"\n\n");
289 prompt.push_str("The `list_sort` stage needs `Record{items, key, descending}` but the pipeline provides a bare `List`.\n");
290 prompt.push_str("**Use config** to supply the parameter fields:\n\n");
291 prompt.push_str("```json\n");
292 prompt.push_str("{\n");
293 prompt.push_str(" \"description\": \"Sort by score and take top 3\",\n");
294 prompt.push_str(" \"version\": \"0.1.0\",\n");
295 prompt.push_str(" \"root\": {\n");
296 prompt.push_str(" \"op\": \"Sequential\",\n");
297 prompt.push_str(" \"stages\": [\n");
298 prompt.push_str(&format!(
299 " {{\"op\": \"Stage\", \"id\": \"{sort_id}\", \"config\": {{\"key\": \"score\", \"descending\": true}}}},\n"
300 ));
301 prompt.push_str(&format!(
302 " {{\"op\": \"Stage\", \"id\": \"{take_id}\", \"config\": {{\"count\": 3}}}},\n"
303 ));
304 prompt.push_str(&format!(
305 " {{\"op\": \"Stage\", \"id\": \"{json_ser_id}\"}}\n"
306 ));
307 prompt.push_str(" ]\n");
308 prompt.push_str(" }\n");
309 prompt.push_str("}\n");
310 prompt.push_str("```\n\n");
311 prompt.push_str("The executor merges config with pipeline data automatically. No Parallel+Const needed.\n\n");
312 prompt.push_str("**When to synthesize instead:** when the operation has complex custom logic (API calls, data transformations that no existing stage covers).\n\n");
313
314 prompt.push_str("## Available Stages\n\n");
316 prompt.push_str("Stages are listed by relevance to your problem (highest first).\n\n");
317
318 for (result, stage) in candidates {
319 prompt.push_str(&format!(
320 "### `{}` — {} _(relevance: {:.2})_\n",
321 stage.id.0, stage.description, result.score
322 ));
323 prompt.push_str(&format!(
324 "- **Input**: `{}`\n- **Output**: `{}`\n",
325 stage.signature.input, stage.signature.output,
326 ));
327
328 for ex in stage.examples.iter().take(2) {
330 let input_str = serde_json::to_string(&ex.input).unwrap_or_default();
331 let output_str = serde_json::to_string(&ex.output).unwrap_or_default();
332 prompt.push_str(&format!("- Example: `{input_str}` → `{output_str}`\n"));
333 }
334 prompt.push('\n');
335 }
336
337 prompt.push_str("## Your Response\n\n");
339 prompt.push_str("Respond with ONLY this JSON (no other text):\n");
340 prompt.push_str("```json\n");
341 prompt.push_str("{\n");
342 prompt.push_str(" \"description\": \"<what this composition does>\",\n");
343 prompt.push_str(" \"version\": \"0.1.0\",\n");
344 prompt.push_str(" \"root\": { <composition using operators above> }\n");
345 prompt.push_str("}\n");
346 prompt.push_str("```\n");
347
348 prompt
349}
350
351fn find_candidate_id(candidates: &[(&SearchResult, &Stage)], needle: &str) -> String {
355 candidates
356 .iter()
357 .find(|(_, s)| s.description.contains(needle))
358 .map(|(_, s)| s.id.0.clone())
359 .unwrap_or_else(|| format!("<{needle}>"))
360}
361
362pub fn build_effect_inference_prompt(code: &str, language: &str) -> String {
367 let mut p = String::new();
368 p.push_str("You are analyzing code to determine its computational effects for the Noether platform.\n\n");
369 p.push_str("## Noether Effect Types\n\n");
370 p.push_str("- **Pure**: No side effects. Same inputs always produce same outputs. No I/O, no randomness.\n");
371 p.push_str("- **Fallible**: The operation may fail or raise an exception.\n");
372 p.push_str("- **Network**: Makes HTTP/TCP/DNS requests or any network I/O.\n");
373 p.push_str("- **NonDeterministic**: Output may vary even with identical inputs (random, timestamp, etc.).\n");
374 p.push_str("- **Llm**: Calls an LLM or AI model API.\n");
375 p.push_str("- **Unknown**: Cannot determine effects from code inspection.\n\n");
376
377 p.push_str(&format!("## Code to Analyze ({language})\n\n"));
378 p.push_str("```\n");
379 p.push_str(code);
380 p.push_str("\n```\n\n");
381
382 p.push_str("## Task\n\n");
383 p.push_str("List ONLY the effects that apply to this code. If the code has no side effects and is deterministic, return `[\"Pure\"]`.\n\n");
384 p.push_str("Rules:\n");
385 p.push_str("- Pure and NonDeterministic are mutually exclusive (non-deterministic implies NOT Pure).\n");
386 p.push_str(
387 "- If the code imports urllib, requests, httpx, aiohttp, or any HTTP library → Network.\n",
388 );
389 p.push_str("- If the code has try/except or can raise → Fallible.\n");
390 p.push_str("- If you cannot determine the effects → Unknown (not Pure).\n\n");
391
392 p.push_str("## Response Format\n\n");
393 p.push_str("Respond with ONLY a JSON array of effect names (no other text):\n");
394 p.push_str("```json\n");
395 p.push_str("[\"Effect1\", \"Effect2\"]\n");
396 p.push_str("```\n");
397 p
398}
399
400pub fn extract_effect_response(response: &str) -> noether_core::effects::EffectSet {
405 use noether_core::effects::{Effect, EffectSet};
406
407 let json_str = match extract_json_array(response) {
408 Some(s) => s,
409 None => return EffectSet::unknown(),
410 };
411
412 let names: Vec<String> = match serde_json::from_str(json_str) {
413 Ok(v) => v,
414 Err(_) => return EffectSet::unknown(),
415 };
416
417 let effects: Vec<Effect> = names
418 .iter()
419 .filter_map(|name| match name.as_str() {
420 "Pure" => Some(Effect::Pure),
421 "Fallible" => Some(Effect::Fallible),
422 "Network" => Some(Effect::Network),
423 "NonDeterministic" => Some(Effect::NonDeterministic),
424 "Llm" => Some(Effect::Llm {
425 model: "unknown".into(),
426 }),
427 "Unknown" => Some(Effect::Unknown),
428 _ => None,
429 })
430 .collect();
431
432 if effects.is_empty() {
433 EffectSet::unknown()
434 } else {
435 EffectSet::new(effects)
436 }
437}
438
439fn extract_json_array(response: &str) -> Option<&str> {
441 if let Some(start) = response.find("```json") {
443 let content = &response[start + 7..];
444 if let Some(end) = content.find("```") {
445 return Some(content[..end].trim());
446 }
447 }
448 if let Some(start) = response.find("```") {
450 let content = &response[start + 3..];
451 if let Some(end) = content.find("```") {
452 let candidate = content[..end].trim();
453 if candidate.starts_with('[') {
454 return Some(candidate);
455 }
456 }
457 }
458 if let Some(start) = response.find('[') {
460 let bytes = response.as_bytes();
461 let mut depth: i32 = 0;
462 let mut in_string = false;
463 let mut escape = false;
464 for (i, &b) in bytes[start..].iter().enumerate() {
465 if escape {
466 escape = false;
467 continue;
468 }
469 if in_string {
470 match b {
471 b'\\' => escape = true,
472 b'"' => in_string = false,
473 _ => {}
474 }
475 continue;
476 }
477 match b {
478 b'"' => in_string = true,
479 b'[' => depth += 1,
480 b']' => {
481 depth -= 1;
482 if depth == 0 {
483 return Some(response[start..start + i + 1].trim());
484 }
485 }
486 _ => {}
487 }
488 }
489 }
490 None
491}
492
493pub fn build_synthesis_prompt(spec: &SynthesisSpec) -> String {
495 let mut p = String::new();
496 p.push_str(
497 "You are generating a stage implementation for the Noether composition platform.\n\n",
498 );
499 p.push_str("## Stage Specification\n\n");
500 p.push_str(&format!("- **Name**: `{}`\n", spec.name));
501 p.push_str(&format!("- **Description**: {}\n", spec.description));
502 p.push_str(&format!("- **Input type**: `{}`\n", spec.input));
503 p.push_str(&format!("- **Output type**: `{}`\n\n", spec.output));
504
505 p.push_str("## Your Task\n\n");
506 p.push_str(
507 "1. Produce at least 3 concrete input/output example pairs matching the type signature.\n",
508 );
509 p.push_str("2. Write a Python function `execute(input_value)` that implements this stage.\n");
510 p.push_str(
511 " `input_value` is a Python dict/str/number/list/bool/None matching the input type.\n",
512 );
513 p.push_str(" Return a value matching the output type.\n\n");
514 p.push_str("## Python Implementation Rules\n\n");
515 p.push_str("- **Prefer Python stdlib over third-party packages** when possible.\n");
516 p.push_str(
517 " - For HTTP: use `urllib.request` / `urllib.parse` (always available), NOT `requests`.\n",
518 );
519 p.push_str(" - For JSON: use `json` (always available).\n");
520 p.push_str(" - For dates: use `datetime` (always available).\n");
521 p.push_str(" - For regex: use `re` (always available).\n");
522 p.push_str("- Only use third-party packages (`requests`, `pandas`, etc.) when there is no stdlib alternative.\n");
523 p.push_str(
524 "- **CRITICAL**: ALL imports MUST be placed at the top of the `execute` function body,\n",
525 );
526 p.push_str(
527 " BEFORE any use of those modules. Never use a module without importing it first.\n\n",
528 );
529 p.push_str("## Correct HTTP Implementation Pattern\n\n");
530 p.push_str("```python\n");
531 p.push_str("def execute(input_value):\n");
532 p.push_str(" # ALWAYS import at the top of execute\n");
533 p.push_str(" import urllib.request, urllib.parse, json\n");
534 p.push_str(" url = 'https://api.example.com/search?' + urllib.parse.urlencode({'q': input_value['query']})\n");
535 p.push_str(" with urllib.request.urlopen(url) as resp:\n");
536 p.push_str(" data = json.loads(resp.read().decode())\n");
537 p.push_str(" return data['items']\n");
538 p.push_str("```\n\n");
539
540 p.push_str("## Output Format\n\n");
541 p.push_str("Respond with ONLY this JSON (no other text):\n");
542 p.push_str("```json\n");
543 p.push_str("{\n");
544 p.push_str(" \"examples\": [\n");
545 p.push_str(" {\"input\": <value>, \"output\": <value>},\n");
546 p.push_str(" {\"input\": <value>, \"output\": <value>},\n");
547 p.push_str(" {\"input\": <value>, \"output\": <value>}\n");
548 p.push_str(" ],\n");
549 p.push_str(" \"implementation\": \"def execute(input_value):\\n ...\",\n");
550 p.push_str(" \"language\": \"python\"\n");
551 p.push_str("}\n");
552 p.push_str("```\n");
553 p
554}
555
556pub fn extract_synthesis_spec(response: &str) -> Option<SynthesisSpec> {
559 let json_str = extract_json(response)?;
560 let v: serde_json::Value = serde_json::from_str(json_str).ok()?;
561 if v.get("action").and_then(|a| a.as_str()) != Some("synthesize") {
562 return None;
563 }
564 let spec = v.get("spec")?;
565 serde_json::from_value(spec.clone()).ok()
566}
567
568pub fn extract_synthesis_response(response: &str) -> Option<SynthesisResponse> {
570 let json_str = extract_json(response)?;
571 serde_json::from_str(json_str).ok()
572}
573
574pub fn extract_json(response: &str) -> Option<&str> {
575 if let Some(start) = response.find("```json") {
577 let json_start = start + 7;
578 let json_content = &response[json_start..];
579 if let Some(end) = json_content.find("```") {
580 return Some(json_content[..end].trim());
581 }
582 }
583
584 if let Some(start) = response.find("```") {
586 let content_start = start + 3;
587 let content = &response[content_start..];
588 let (skip, rest) = match content.find('\n') {
590 Some(nl) => {
591 let first_line = content[..nl].trim();
592 if first_line.starts_with('{') {
593 (0, content)
594 } else {
595 (nl + 1, &content[nl + 1..])
596 }
597 }
598 None => (0, content),
599 };
600 let _ = skip;
601 if let Some(end) = rest.find("```") {
602 let candidate = rest[..end].trim();
603 if candidate.starts_with('{') {
604 return Some(candidate);
605 }
606 }
607 }
608
609 if let Some(brace_start) = response.find('{') {
612 let bytes = response.as_bytes();
613 let mut depth: i32 = 0;
614 let mut in_string = false;
615 let mut escape = false;
616 let mut brace_end: Option<usize> = None;
617
618 for (i, &b) in bytes[brace_start..].iter().enumerate() {
619 if escape {
620 escape = false;
621 continue;
622 }
623 if in_string {
624 match b {
625 b'\\' => escape = true,
626 b'"' => in_string = false,
627 _ => {}
628 }
629 continue;
630 }
631 match b {
632 b'"' => in_string = true,
633 b'{' => depth += 1,
634 b'}' => {
635 depth -= 1;
636 if depth == 0 {
637 brace_end = Some(brace_start + i + 1);
638 break;
639 }
640 }
641 _ => {}
642 }
643 }
644
645 if let Some(end) = brace_end {
646 let candidate = response[brace_start..end].trim();
647 if !candidate.is_empty() {
648 return Some(candidate);
649 }
650 }
651 }
652
653 None
654}
655
656#[cfg(test)]
657mod tests {
658 use super::*;
659 use crate::index::SearchResult;
660 use noether_core::stage::StageId;
661
662 fn make_search_result(id: &str, score: f32) -> SearchResult {
663 SearchResult {
664 stage_id: StageId(id.into()),
665 score,
666 signature_score: score,
667 semantic_score: score,
668 example_score: score,
669 }
670 }
671
672 #[test]
673 fn extract_json_from_code_block() {
674 let response = "Here's the graph:\n```json\n{\"test\": true}\n```\nDone.";
675 assert_eq!(extract_json(response), Some("{\"test\": true}"));
676 }
677
678 #[test]
679 fn extract_json_from_plain_block() {
680 let response = "```\n{\"test\": true}\n```";
681 assert_eq!(extract_json(response), Some("{\"test\": true}"));
682 }
683
684 #[test]
685 fn extract_raw_json() {
686 let response = "{\"test\": true}";
687 assert_eq!(extract_json(response), Some("{\"test\": true}"));
688 }
689
690 #[test]
691 fn extract_json_none_for_text() {
692 let response = "No JSON here, just text.";
693 assert_eq!(extract_json(response), None);
694 }
695
696 #[test]
697 fn extract_json_with_whitespace() {
698 let response = " \n```json\n {\"a\": 1} \n```\n ";
699 assert_eq!(extract_json(response), Some("{\"a\": 1}"));
700 }
701
702 #[test]
703 fn extract_synthesis_spec_parses_valid_request() {
704 let input_json = serde_json::to_string(&NType::Text).unwrap();
705 let output_json = serde_json::to_string(&NType::Number).unwrap();
706 let response = format!(
707 "```json\n{}\n```",
708 serde_json::json!({
709 "action": "synthesize",
710 "spec": {
711 "name": "count_words",
712 "description": "Count the number of words in a text",
713 "input": serde_json::from_str::<serde_json::Value>(&input_json).unwrap(),
714 "output": serde_json::from_str::<serde_json::Value>(&output_json).unwrap(),
715 "rationale": "No existing stage counts words"
716 }
717 })
718 );
719 let spec = extract_synthesis_spec(&response).unwrap();
720 assert_eq!(spec.name, "count_words");
721 assert_eq!(spec.input, NType::Text);
722 assert_eq!(spec.output, NType::Number);
723 }
724
725 #[test]
726 fn extract_synthesis_spec_returns_none_for_composition_graph() {
727 let response = "```json\n{\"description\":\"test\",\"version\":\"0.1.0\",\"root\":{\"op\":\"Stage\",\"id\":\"abc\"}}\n```";
728 assert!(extract_synthesis_spec(response).is_none());
729 }
730
731 #[test]
732 fn extract_synthesis_response_parses_examples_and_code() {
733 let response = "```json\n{\"examples\":[{\"input\":\"hello world\",\"output\":2},{\"input\":\"foo\",\"output\":1}],\"implementation\":\"def execute(v): return len(v.split())\",\"language\":\"python\"}\n```";
734 let resp = extract_synthesis_response(response).unwrap();
735 assert_eq!(resp.examples.len(), 2);
736 assert_eq!(resp.language, "python");
737 assert!(resp.implementation.contains("execute"));
738 }
739
740 #[test]
741 fn build_synthesis_prompt_contains_spec_fields() {
742 let spec = SynthesisSpec {
743 name: "reverse_text".into(),
744 description: "Reverse a string".into(),
745 input: NType::Text,
746 output: NType::Text,
747 rationale: "no existing stage reverses text".into(),
748 };
749 let prompt = build_synthesis_prompt(&spec);
750 assert!(prompt.contains("reverse_text"));
751 assert!(prompt.contains("Reverse a string"));
752 assert!(prompt.contains("execute(input_value)"));
753 }
754
755 #[test]
756 fn few_shot_uses_real_ids_when_candidates_present() {
757 use noether_core::stdlib::load_stdlib;
758
759 let stages = load_stdlib();
760 let parse_json = stages
761 .iter()
762 .find(|s| s.description.contains("Parse a JSON string"))
763 .unwrap();
764 let to_json = stages
765 .iter()
766 .find(|s| s.description.contains("Serialize any value to a JSON"))
767 .unwrap();
768
769 let r1 = make_search_result(&parse_json.id.0, 0.9);
770 let r2 = make_search_result(&to_json.id.0, 0.8);
771 let candidates: Vec<(&SearchResult, &Stage)> = vec![(&r1, parse_json), (&r2, to_json)];
772
773 let prompt = build_system_prompt(&candidates);
774
775 assert!(
777 prompt.contains(&parse_json.id.0),
778 "prompt should contain real parse_json hash"
779 );
780 assert!(
781 prompt.contains(&to_json.id.0),
782 "prompt should contain real to_json hash"
783 );
784 }
785
786 #[test]
787 fn few_shot_falls_back_to_placeholder_when_stages_absent() {
788 let prompt = build_system_prompt(&[]);
789 assert!(
791 prompt.contains("<Parse a JSON string>"),
792 "expected placeholder when parse_json not in candidates"
793 );
794 }
795
796 #[test]
797 fn prompt_contains_branch_guidance() {
798 let prompt = build_system_prompt(&[]);
799 assert!(
800 prompt.contains("predicate"),
801 "prompt should explain Branch predicate"
802 );
803 assert!(
804 prompt.contains("original input"),
805 "prompt should clarify that if_true/if_false receive original input"
806 );
807 assert!(
808 prompt.contains("Stage Config"),
809 "prompt should have Stage Config section"
810 );
811 assert!(
812 prompt.contains("\"Const\""),
813 "prompt should list Const as a valid op"
814 );
815 assert!(
816 prompt.contains("config") && prompt.contains("key"),
817 "prompt should explain config pattern for parameterized stages"
818 );
819 }
820
821 #[test]
822 fn candidates_show_relevance_score() {
823 use noether_core::stdlib::load_stdlib;
824
825 let stages = load_stdlib();
826 let stage = stages.first().unwrap();
827 let r = make_search_result(&stage.id.0, 0.75);
828 let candidates: Vec<(&SearchResult, &Stage)> = vec![(&r, stage)];
829
830 let prompt = build_system_prompt(&candidates);
831 assert!(
832 prompt.contains("relevance: 0.75"),
833 "prompt should display the fused relevance score"
834 );
835 }
836}