chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,15 @@
|
||||
"""SkillOpt Gradient -- trajectory analysis and patch generation.
|
||||
|
||||
Analogous to gradient computation in neural network training: analyzes
|
||||
minibatch rollout trajectories to produce skill-edit patches (the "gradient"
|
||||
that drives skill updates).
|
||||
|
||||
Modules
|
||||
-------
|
||||
- reflect: minibatch trajectory analysis (gradient computation)
|
||||
- aggregate: hierarchical patch merging (gradient aggregation)
|
||||
"""
|
||||
from skillopt.gradient.reflect import ( # noqa: F401
|
||||
run_minibatch_reflect,
|
||||
)
|
||||
from skillopt.gradient.aggregate import merge_patches # noqa: F401
|
||||
@@ -0,0 +1,253 @@
|
||||
"""ReflACT Aggregate stage — hierarchical patch merging.
|
||||
|
||||
The Aggregate stage takes independently-generated patches from the Reflect
|
||||
stage and merges them into a single coherent patch via hierarchical LLM calls.
|
||||
Failure-driven patches take priority over success-driven ones.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from skillopt.model import chat_optimizer
|
||||
from skillopt.optimizer.meta_skill import format_meta_skill_context
|
||||
from skillopt.optimizer.update_modes import (
|
||||
get_payload_items,
|
||||
is_full_rewrite_minibatch_mode,
|
||||
is_rewrite_mode,
|
||||
normalize_update_mode,
|
||||
payload_key,
|
||||
payload_label,
|
||||
)
|
||||
from skillopt.prompts import load_prompt
|
||||
from skillopt.utils import extract_json
|
||||
|
||||
|
||||
# ── Internal helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
def _merge_batch(
|
||||
skill_content: str,
|
||||
patches: list[dict],
|
||||
system_prompt: str,
|
||||
update_mode: str,
|
||||
meta_skill_context: str = "",
|
||||
level: int = 1,
|
||||
) -> dict:
|
||||
"""Call optimizer LLM to merge a batch of patches into one."""
|
||||
patches_text = json.dumps(patches, ensure_ascii=False, indent=2)
|
||||
user = (
|
||||
f"## Current Skill\n{skill_content}\n\n"
|
||||
f"## Patches to merge ({len(patches)} total, merge level {level})\n{patches_text}"
|
||||
)
|
||||
optimizer_ctx = format_meta_skill_context(meta_skill_context)
|
||||
if optimizer_ctx:
|
||||
user = f"{optimizer_ctx}\n\n{user}"
|
||||
try:
|
||||
response, _ = chat_optimizer(
|
||||
system=system_prompt,
|
||||
user=user,
|
||||
max_completion_tokens=64000 if is_full_rewrite_minibatch_mode(update_mode) else 16384,
|
||||
retries=3,
|
||||
stage="merge",
|
||||
)
|
||||
merged = extract_json(response)
|
||||
key = payload_key(update_mode)
|
||||
if merged and key in merged:
|
||||
for e in merged.get(key, []):
|
||||
e["merge_level"] = level
|
||||
return merged
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
# Fallback: concatenate all edits
|
||||
all_edits = []
|
||||
for p in patches:
|
||||
for e in get_payload_items(p, update_mode):
|
||||
e.setdefault("merge_level", level)
|
||||
all_edits.append(e)
|
||||
return {"reasoning": "fallback concatenation", payload_key(update_mode): all_edits}
|
||||
|
||||
|
||||
def _hierarchical_merge(
|
||||
skill_content: str,
|
||||
patches: list[dict],
|
||||
system_prompt: str,
|
||||
update_mode: str,
|
||||
batch_size: int,
|
||||
verbose: bool,
|
||||
label: str = "",
|
||||
workers: int = 16,
|
||||
meta_skill_context: str = "",
|
||||
) -> dict:
|
||||
"""Hierarchically merge N patches using the given system prompt.
|
||||
|
||||
Same-level batches are executed in PARALLEL via ThreadPoolExecutor.
|
||||
"""
|
||||
if not patches:
|
||||
return {"reasoning": "no patches", payload_key(update_mode): []}
|
||||
if len(patches) == 1:
|
||||
return patches[0]
|
||||
|
||||
current = list(patches)
|
||||
level = 0
|
||||
while len(current) > 1:
|
||||
level += 1
|
||||
batches: list[tuple[int, list[dict]]] = []
|
||||
for i in range(0, len(current), batch_size):
|
||||
batch = current[i : i + batch_size]
|
||||
batches.append((i, batch))
|
||||
|
||||
if verbose:
|
||||
print(
|
||||
f" [aggregate {label}] level={level} "
|
||||
f"{len(current)} patches → {len(batches)} batches "
|
||||
f"(parallel, batch_size={batch_size})"
|
||||
)
|
||||
|
||||
next_level: list[dict | None] = [None] * len(batches)
|
||||
|
||||
to_merge: list[tuple[int, list[dict]]] = []
|
||||
for idx, (i, batch) in enumerate(batches):
|
||||
if len(batch) == 1:
|
||||
next_level[idx] = batch[0]
|
||||
else:
|
||||
to_merge.append((idx, batch))
|
||||
|
||||
if to_merge:
|
||||
with ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
futs = {
|
||||
ex.submit(
|
||||
_merge_batch, skill_content, batch, system_prompt, update_mode,
|
||||
meta_skill_context, level,
|
||||
): idx
|
||||
for idx, batch in to_merge
|
||||
}
|
||||
for fut in as_completed(futs):
|
||||
idx = futs[fut]
|
||||
next_level[idx] = fut.result()
|
||||
if verbose:
|
||||
batch_i, batch_data = batches[idx]
|
||||
n_edits = len(get_payload_items(next_level[idx], update_mode))
|
||||
print(
|
||||
f" [aggregate {label}] level={level} "
|
||||
f"batch [{batch_i}:{batch_i+len(batch_data)}] "
|
||||
f"→ 1 patch ({n_edits} {payload_label(update_mode)})"
|
||||
)
|
||||
|
||||
current = [x for x in next_level if x is not None]
|
||||
|
||||
return current[0]
|
||||
|
||||
|
||||
# ── Public API ────────────────────────────────────────────────────────────────
|
||||
|
||||
def merge_patches(
|
||||
skill_content: str,
|
||||
failure_patches: list[dict],
|
||||
success_patches: list[dict],
|
||||
batch_size: int = 8,
|
||||
verbose: bool = True,
|
||||
workers: int = 16,
|
||||
update_mode: str = "patch",
|
||||
meta_skill_context: str = "",
|
||||
) -> dict:
|
||||
"""Failure-first hierarchical merge with support count tracking.
|
||||
|
||||
1. Merge failure patches independently (parallel)
|
||||
2. Merge success patches independently (parallel)
|
||||
3. Final merge: combine both groups with failure priority
|
||||
|
||||
Returns a merged :class:`~skillopt.types.Patch` dict (``edits`` + ``reasoning``).
|
||||
"""
|
||||
if verbose:
|
||||
print(
|
||||
f" [3/6 AGGREGATE] "
|
||||
f"failure={len(failure_patches)} success={len(success_patches)} "
|
||||
f"(parallel, workers={workers})"
|
||||
)
|
||||
|
||||
update_mode = normalize_update_mode(update_mode)
|
||||
if is_full_rewrite_minibatch_mode(update_mode):
|
||||
merge_failure_prompt = load_prompt("merge_failure_full_rewrite")
|
||||
merge_success_prompt = load_prompt("merge_success_full_rewrite")
|
||||
merge_final_prompt = load_prompt("merge_final_full_rewrite")
|
||||
elif is_rewrite_mode(update_mode):
|
||||
merge_failure_prompt = load_prompt("merge_failure_rewrite")
|
||||
merge_success_prompt = load_prompt("merge_success_rewrite")
|
||||
merge_final_prompt = load_prompt("merge_final_rewrite")
|
||||
else:
|
||||
merge_failure_prompt = load_prompt("merge_failure")
|
||||
merge_success_prompt = load_prompt("merge_success")
|
||||
merge_final_prompt = load_prompt("merge_final")
|
||||
|
||||
failure_merged = _hierarchical_merge(
|
||||
skill_content, failure_patches, merge_failure_prompt, update_mode,
|
||||
batch_size, verbose, label="failure", workers=workers,
|
||||
meta_skill_context=meta_skill_context,
|
||||
)
|
||||
|
||||
success_merged = _hierarchical_merge(
|
||||
skill_content, success_patches, merge_success_prompt, update_mode,
|
||||
batch_size, verbose, label="success", workers=workers,
|
||||
meta_skill_context=meta_skill_context,
|
||||
)
|
||||
|
||||
f_edits = get_payload_items(failure_merged, update_mode)
|
||||
s_edits = get_payload_items(success_merged, update_mode)
|
||||
|
||||
if not f_edits and not s_edits:
|
||||
return {"reasoning": "no updates from either group", payload_key(update_mode): []}
|
||||
if not s_edits:
|
||||
return failure_merged
|
||||
if not f_edits:
|
||||
return success_merged
|
||||
|
||||
combined_patches = [failure_merged, success_merged]
|
||||
combined_text = json.dumps(combined_patches, ensure_ascii=False, indent=2)
|
||||
if is_full_rewrite_minibatch_mode(update_mode):
|
||||
item_label = payload_label(update_mode)
|
||||
user = (
|
||||
f"## Current Skill\n{skill_content}\n\n"
|
||||
f"## Two pre-merged candidate groups to combine\n"
|
||||
f"Group 1 (from failed trajectories): "
|
||||
f"{len(f_edits)} {item_label}\n"
|
||||
f"Group 2 (from successful trajectories): "
|
||||
f"{len(s_edits)} {item_label}\n\n"
|
||||
f"{combined_text}"
|
||||
)
|
||||
else:
|
||||
user = (
|
||||
f"## Current Skill\n{skill_content}\n\n"
|
||||
f"## Two pre-merged patch groups to combine\n"
|
||||
f"Group 1 (failure-driven, HIGH priority): "
|
||||
f"{len(f_edits)} edits\n"
|
||||
f"Group 2 (success-driven, lower priority): "
|
||||
f"{len(s_edits)} edits\n\n"
|
||||
f"{combined_text}"
|
||||
)
|
||||
optimizer_ctx = format_meta_skill_context(meta_skill_context)
|
||||
if optimizer_ctx:
|
||||
user = f"{optimizer_ctx}\n\n{user}"
|
||||
try:
|
||||
response, _ = chat_optimizer(
|
||||
system=merge_final_prompt,
|
||||
user=user,
|
||||
max_completion_tokens=64000 if is_full_rewrite_minibatch_mode(update_mode) else 16384,
|
||||
retries=3,
|
||||
stage="merge",
|
||||
)
|
||||
final = extract_json(response)
|
||||
key = payload_key(update_mode)
|
||||
if final and key in final:
|
||||
if verbose:
|
||||
print(
|
||||
f" [aggregate final] "
|
||||
f"{len(f_edits)}+{len(s_edits)} → {len(final[key])} {payload_label(update_mode)}"
|
||||
)
|
||||
return final
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
return {
|
||||
"reasoning": "fallback: failure first, then success",
|
||||
payload_key(update_mode): f_edits + s_edits,
|
||||
}
|
||||
@@ -0,0 +1,635 @@
|
||||
"""ReflACT core Reflect engine -- minibatch trajectory analysis.
|
||||
|
||||
Provides environment-agnostic minibatch trajectory analysis: instead of
|
||||
analyzing each trajectory independently, trajectories are grouped into
|
||||
minibatches of size M and analyzed together -- analogous to minibatch SGD
|
||||
vs per-sample SGD in neural network training.
|
||||
|
||||
Two-level prompt priority system:
|
||||
|
||||
1. **Custom prompt** (adapter returns non-None) -- used as-is.
|
||||
2. **Generic default prompt** (adapter returns None) -- built-in defaults
|
||||
that work for any environment without configuration.
|
||||
|
||||
Public API
|
||||
----------
|
||||
- :func:`fmt_trajectory` -- format one conversation into text
|
||||
- :func:`fmt_minibatch_trajectories` -- format multiple trajectories for batch analysis
|
||||
- :func:`run_error_analyst_minibatch` -- one optimizer call for a group of failures
|
||||
- :func:`run_success_analyst_minibatch` -- one optimizer call for a group of successes
|
||||
- :func:`run_minibatch_reflect` -- full reflect stage dispatcher
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from skillopt.model import chat_optimizer
|
||||
from skillopt.optimizer.meta_skill import format_meta_skill_context
|
||||
from skillopt.optimizer.skill_aware import (
|
||||
augment_error_prompt,
|
||||
augment_success_prompt,
|
||||
extract_appendix_notes,
|
||||
get_skill_aware_appendix_source,
|
||||
is_skill_aware_enabled,
|
||||
)
|
||||
from skillopt.optimizer.update_modes import (
|
||||
get_payload_items,
|
||||
is_full_rewrite_minibatch_mode,
|
||||
normalize_update_mode,
|
||||
payload_key,
|
||||
payload_label,
|
||||
truncate_payload,
|
||||
)
|
||||
from skillopt.prompts import load_prompt
|
||||
from skillopt.utils import extract_json
|
||||
|
||||
|
||||
# ── Trajectory formatting ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _clip_text(value, limit: int | None = None) -> str:
|
||||
"""Render optional trajectory fields. Truncation is disabled: the optimizer
|
||||
is given the full content so it can see exactly what the agent saw/did.
|
||||
|
||||
``limit`` is accepted for backward compatibility but ignored.
|
||||
"""
|
||||
if value is None:
|
||||
return ""
|
||||
return str(value)
|
||||
|
||||
|
||||
def fmt_trajectory(
|
||||
conversation: list[dict],
|
||||
max_chars: int | None = None,
|
||||
) -> str:
|
||||
"""Format a conversation list into analyst-readable text.
|
||||
|
||||
Accepts two common formats:
|
||||
|
||||
1. Tool-call records: ``{"type": "tool_call", "cmd": ..., "obs": ...}``
|
||||
2. Step records: ``{"step": N, "action": ..., "env_feedback": ..., "reasoning": ...}``
|
||||
|
||||
Any other dict is rendered via its ``"content"`` key.
|
||||
"""
|
||||
lines: list[str] = []
|
||||
for item in conversation:
|
||||
if not isinstance(item, dict):
|
||||
lines.append(f"[agent] {_clip_text(item)}")
|
||||
continue
|
||||
if item.get("type") == "tool_call":
|
||||
cmd = _clip_text(item.get("cmd"))
|
||||
obs = _clip_text(item.get("obs"))
|
||||
lines.append(f"[action] {cmd}")
|
||||
lines.append(f"[obs] {obs}")
|
||||
elif "action" in item and "env_feedback" in item:
|
||||
step = item.get("step", "?")
|
||||
reasoning = _clip_text(item.get("reasoning"))
|
||||
action = _clip_text(item.get("action"))
|
||||
feedback = _clip_text(item.get("env_feedback"))
|
||||
if reasoning:
|
||||
lines.append(f"[step {step} think] {reasoning}")
|
||||
lines.append(f"[step {step} action] {action}")
|
||||
lines.append(f"[step {step} obs] {feedback}")
|
||||
elif item.get("role") == "system":
|
||||
# Post-execution verification / enrichment info
|
||||
msg = _clip_text(item.get("content"))
|
||||
lines.append(f"[verification] {msg}")
|
||||
else:
|
||||
msg = _clip_text(item.get("content"))
|
||||
role = item.get("role", "agent")
|
||||
lines.append(f"[{role}] {msg}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# ── Minibatch trajectory formatting ──────────────────────────────────────────
|
||||
|
||||
|
||||
def fmt_minibatch_trajectories(
|
||||
items: list[dict],
|
||||
prediction_dir: str,
|
||||
) -> str:
|
||||
"""Format multiple trajectories for minibatch analyst consumption.
|
||||
|
||||
Each item is a rollout result dict with ``"id"``, ``"task_description"``,
|
||||
``"task_type"``, ``"fail_reason"``, etc. Reads ``conversation.json``
|
||||
for each and formats them together with trajectory headers.
|
||||
|
||||
If available, includes the spreadsheet preview and target system prompt
|
||||
so the analyst can see what the agent saw.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
items : list[dict]
|
||||
Rollout result dicts belonging to one minibatch.
|
||||
prediction_dir : str
|
||||
Path to ``predictions/`` directory containing per-task
|
||||
``<task_id>/conversation.json`` files.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Formatted text with all trajectories separated by ``---``.
|
||||
"""
|
||||
parts: list[str] = []
|
||||
for idx, item in enumerate(items, 1):
|
||||
tid = str(item["id"])
|
||||
conv_path = os.path.join(prediction_dir, tid, "conversation.json")
|
||||
if not os.path.exists(conv_path):
|
||||
continue
|
||||
with open(conv_path) as f:
|
||||
conversation = json.load(f)
|
||||
if not conversation:
|
||||
continue
|
||||
|
||||
traj_text = fmt_trajectory(conversation)
|
||||
header = (
|
||||
f"### Trajectory {idx} (id={tid})\n"
|
||||
f"Task: {item.get('task_description', item.get('instruction', ''))}\n"
|
||||
f"Task type: {item.get('task_type', item.get('instruction_type', ''))}\n"
|
||||
)
|
||||
fail_reason = item.get("fail_reason", "")
|
||||
if fail_reason:
|
||||
header += f"Failure reason: {fail_reason}\n"
|
||||
header += f"Steps: {item.get('n_turns', '?')}\n"
|
||||
|
||||
reference_text = str(item.get("reference_text") or "").strip()
|
||||
if reference_text:
|
||||
header += (
|
||||
f"\n#### Hidden Reference\n"
|
||||
f"{reference_text}\n"
|
||||
)
|
||||
|
||||
# ── Append target context (what the agent saw) ──────────────
|
||||
target_prompt = item.get("target_system_prompt", "")
|
||||
if not target_prompt:
|
||||
prompt_path = os.path.join(prediction_dir, tid, "target_system_prompt.txt")
|
||||
if os.path.exists(prompt_path):
|
||||
with open(prompt_path) as f:
|
||||
target_prompt = f.read()
|
||||
if target_prompt:
|
||||
header += (
|
||||
f"\n#### Target System Prompt\n"
|
||||
f"{target_prompt}\n"
|
||||
)
|
||||
|
||||
user_prompt = item.get("target_user_prompt", "")
|
||||
if not user_prompt:
|
||||
user_prompt_path = os.path.join(prediction_dir, tid, "target_user_prompt.txt")
|
||||
if os.path.exists(user_prompt_path):
|
||||
with open(user_prompt_path) as f:
|
||||
user_prompt = f.read()
|
||||
if user_prompt:
|
||||
header += (
|
||||
f"\n#### Target User Prompt\n"
|
||||
f"{user_prompt}\n"
|
||||
)
|
||||
|
||||
if os.environ.get("REFLACT_CODEX_TRACE_TO_OPTIMIZER", "0") == "1":
|
||||
codex_trace_summary = item.get("codex_trace_summary", "")
|
||||
if not codex_trace_summary:
|
||||
codex_trace_summary_path = os.path.join(prediction_dir, tid, "codex_trace_summary.txt")
|
||||
if os.path.exists(codex_trace_summary_path):
|
||||
with open(codex_trace_summary_path) as f:
|
||||
codex_trace_summary = f.read()
|
||||
if codex_trace_summary:
|
||||
header += (
|
||||
f"\n#### Codex Trace Summary\n"
|
||||
f"{codex_trace_summary}\n"
|
||||
)
|
||||
|
||||
codex_probe_trace_steps = str(item.get("codex_probe_trace_steps") or "").strip()
|
||||
if codex_probe_trace_steps:
|
||||
header += (
|
||||
f"\n#### Codex Trace Steps\n"
|
||||
f"{codex_probe_trace_steps}\n"
|
||||
)
|
||||
|
||||
preview = item.get("spreadsheet_preview", "")
|
||||
if not preview:
|
||||
preview_path = os.path.join(prediction_dir, tid, "spreadsheet_preview.txt")
|
||||
if os.path.exists(preview_path):
|
||||
with open(preview_path) as f:
|
||||
preview = f.read()
|
||||
if preview:
|
||||
header += (
|
||||
f"\n#### Spreadsheet Preview\n"
|
||||
f"{preview}\n"
|
||||
)
|
||||
|
||||
parts.append(header + "\n" + traj_text)
|
||||
|
||||
return "\n\n---\n\n".join(parts)
|
||||
|
||||
|
||||
# ── Prompt resolution ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _resolve_prompt(custom: str | None, default_name: str, update_mode: str = "patch") -> str:
|
||||
"""Return *custom* if provided (non-None), otherwise load from file."""
|
||||
if custom is not None:
|
||||
return custom
|
||||
mode = normalize_update_mode(update_mode)
|
||||
actual_name = default_name
|
||||
if is_full_rewrite_minibatch_mode(mode):
|
||||
full_name = f"{default_name}_full_rewrite"
|
||||
try:
|
||||
return load_prompt(full_name)
|
||||
except FileNotFoundError:
|
||||
actual_name = default_name
|
||||
elif mode == "rewrite_from_suggestions":
|
||||
rewrite_name = f"{default_name}_rewrite"
|
||||
try:
|
||||
return load_prompt(rewrite_name)
|
||||
except FileNotFoundError:
|
||||
actual_name = default_name
|
||||
return load_prompt(actual_name)
|
||||
|
||||
|
||||
# ── Minibatch analysts ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def run_error_analyst_minibatch(
|
||||
skill_content: str,
|
||||
items: list[dict],
|
||||
prediction_dir: str,
|
||||
edit_budget: int = 4,
|
||||
*,
|
||||
system_prompt: str | None = None,
|
||||
rejection_context: str = "",
|
||||
trajectory_memory_context: str = "",
|
||||
step_buffer_context: str = "",
|
||||
meta_skill_context: str = "",
|
||||
update_mode: str = "patch",
|
||||
skill_aware_reflection: bool = False,
|
||||
) -> dict | None:
|
||||
"""Analyze a minibatch of failed trajectories in one optimizer call.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
skill_content : str
|
||||
Current skill document text.
|
||||
items : list[dict]
|
||||
Rollout result dicts (all should have ``hard=0``).
|
||||
prediction_dir : str
|
||||
Path to ``predictions/`` directory.
|
||||
edit_budget : int
|
||||
Maximum number of edits (L).
|
||||
system_prompt : str | None
|
||||
Custom system prompt. ``None`` = use generic default.
|
||||
rejection_context : str
|
||||
*Deprecated* — use ``step_buffer_context``.
|
||||
trajectory_memory_context : str
|
||||
*Deprecated* — use ``step_buffer_context``.
|
||||
step_buffer_context : str
|
||||
Unified summary of previous steps (failure patterns + rejected edits).
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict | None
|
||||
Patch dict with ``source_type="failure"``, or ``None`` on error.
|
||||
"""
|
||||
mode = normalize_update_mode(update_mode)
|
||||
actual_system = _resolve_prompt(system_prompt, "analyst_error", mode)
|
||||
# Skill-aware reflection: augment the resolved prompt at runtime so both
|
||||
# env-specific and generic analyst prompts get the defect/lapse instruction.
|
||||
# When the toggle is off this is a no-op (prompt byte-identical to baseline).
|
||||
if skill_aware_reflection and not is_full_rewrite_minibatch_mode(mode):
|
||||
actual_system = augment_error_prompt(actual_system)
|
||||
|
||||
trajectories_text = fmt_minibatch_trajectories(items, prediction_dir)
|
||||
if not trajectories_text.strip():
|
||||
return None
|
||||
|
||||
user = (
|
||||
f"## Current Skill\n{skill_content}\n\n"
|
||||
)
|
||||
if is_full_rewrite_minibatch_mode(mode):
|
||||
user += (
|
||||
f"## Update Format\n"
|
||||
f"Produce one complete replacement skill candidate for this minibatch. "
|
||||
f"Do not output edits, patches, or revise suggestions.\n\n"
|
||||
)
|
||||
else:
|
||||
user += (
|
||||
f"## {payload_label(mode, title=True)} Budget\n"
|
||||
f"Produce at most L={edit_budget} {payload_label(mode)}.\n\n"
|
||||
)
|
||||
# Unified step buffer context (preferred)
|
||||
ctx = step_buffer_context or rejection_context or ""
|
||||
if trajectory_memory_context:
|
||||
ctx = f"{ctx}\n{trajectory_memory_context}" if ctx else trajectory_memory_context
|
||||
if ctx.strip():
|
||||
user += f"## Previous Steps in This Epoch\n{ctx}\n\n"
|
||||
optimizer_ctx = format_meta_skill_context(meta_skill_context)
|
||||
if optimizer_ctx:
|
||||
user += optimizer_ctx + "\n\n"
|
||||
user += f"## Failed Trajectories ({len(items)} total)\n{trajectories_text}"
|
||||
|
||||
try:
|
||||
response, _ = chat_optimizer(
|
||||
system=actual_system, user=user,
|
||||
max_completion_tokens=64000 if is_full_rewrite_minibatch_mode(mode) else 16384,
|
||||
retries=3,
|
||||
stage="analyst",
|
||||
)
|
||||
result = extract_json(response)
|
||||
if not result:
|
||||
return None
|
||||
notes = extract_appendix_notes(result) if skill_aware_reflection else []
|
||||
if "patch" in result:
|
||||
result["source_type"] = "failure"
|
||||
if not is_full_rewrite_minibatch_mode(mode):
|
||||
truncate_payload(result["patch"], edit_budget, mode)
|
||||
if skill_aware_reflection:
|
||||
result["appendix_notes"] = notes
|
||||
return result
|
||||
# Skill-aware: a batch may legitimately yield ONLY execution-lapse notes
|
||||
# (no body edit). Return a no-op patch so the notes still reach the
|
||||
# trainer via all_raw_patches; empty edits are dropped from the body
|
||||
# pipeline by _normalise_patches, so body behavior is unchanged.
|
||||
if skill_aware_reflection and notes:
|
||||
return {
|
||||
"source_type": "failure",
|
||||
"patch": {"reasoning": "execution-lapse only", "edits": []},
|
||||
"appendix_notes": notes,
|
||||
}
|
||||
except Exception: # noqa: BLE001
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
|
||||
def run_success_analyst_minibatch(
|
||||
skill_content: str,
|
||||
items: list[dict],
|
||||
prediction_dir: str,
|
||||
edit_budget: int = 4,
|
||||
*,
|
||||
system_prompt: str | None = None,
|
||||
trajectory_memory_context: str = "",
|
||||
step_buffer_context: str = "",
|
||||
meta_skill_context: str = "",
|
||||
update_mode: str = "patch",
|
||||
skill_aware_reflection: bool = False,
|
||||
emit_appendix_notes: bool = True,
|
||||
) -> dict | None:
|
||||
"""Analyze a minibatch of successful trajectories in one optimizer call.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
system_prompt : str | None
|
||||
Custom system prompt. ``None`` = use generic default.
|
||||
trajectory_memory_context : str
|
||||
*Deprecated* — use ``step_buffer_context``.
|
||||
step_buffer_context : str
|
||||
Unified summary of previous steps (failure patterns + rejected edits).
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict | None
|
||||
Patch dict with ``source_type="success"``, or ``None`` on error.
|
||||
"""
|
||||
mode = normalize_update_mode(update_mode)
|
||||
actual_system = _resolve_prompt(system_prompt, "analyst_success", mode)
|
||||
# Only augment + parse appendix notes on the success side when allowed.
|
||||
# failure_only mode (paper-faithful S_app) suppresses success-side notes.
|
||||
sa_emit = skill_aware_reflection and emit_appendix_notes
|
||||
if sa_emit and not is_full_rewrite_minibatch_mode(mode):
|
||||
actual_system = augment_success_prompt(actual_system)
|
||||
|
||||
trajectories_text = fmt_minibatch_trajectories(items, prediction_dir)
|
||||
if not trajectories_text.strip():
|
||||
return None
|
||||
|
||||
user = (
|
||||
f"## Current Skill\n{skill_content}\n\n"
|
||||
)
|
||||
if is_full_rewrite_minibatch_mode(mode):
|
||||
user += (
|
||||
f"## Update Format\n"
|
||||
f"Produce one complete replacement skill candidate for this minibatch. "
|
||||
f"Do not output edits, patches, or revise suggestions.\n\n"
|
||||
)
|
||||
else:
|
||||
user += (
|
||||
f"## {payload_label(mode, title=True)} Budget\n"
|
||||
f"Produce at most L={edit_budget} {payload_label(mode)}.\n\n"
|
||||
)
|
||||
ctx = step_buffer_context or trajectory_memory_context or ""
|
||||
if ctx.strip():
|
||||
user += f"## Previous Steps in This Epoch\n{ctx}\n\n"
|
||||
optimizer_ctx = format_meta_skill_context(meta_skill_context)
|
||||
if optimizer_ctx:
|
||||
user += optimizer_ctx + "\n\n"
|
||||
user += f"## Successful Trajectories ({len(items)} total)\n{trajectories_text}"
|
||||
|
||||
try:
|
||||
response, _ = chat_optimizer(
|
||||
system=actual_system, user=user,
|
||||
max_completion_tokens=64000 if is_full_rewrite_minibatch_mode(mode) else 16384,
|
||||
retries=3,
|
||||
stage="analyst",
|
||||
)
|
||||
result = extract_json(response)
|
||||
if result and "patch" in result:
|
||||
result["source_type"] = "success"
|
||||
if not is_full_rewrite_minibatch_mode(mode):
|
||||
truncate_payload(result["patch"], edit_budget, mode)
|
||||
if sa_emit:
|
||||
result["appendix_notes"] = extract_appendix_notes(result)
|
||||
return result
|
||||
except Exception: # noqa: BLE001
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
|
||||
# ── Minibatch reflect dispatcher ────────────────────────────────────────────
|
||||
|
||||
|
||||
def _split_minibatches(items: list, batch_size: int) -> list[list]:
|
||||
"""Split items into minibatches of at most *batch_size*."""
|
||||
return [items[i : i + batch_size] for i in range(0, len(items), batch_size)]
|
||||
|
||||
|
||||
def _shuffle_for_minibatch(items: list, seed: int | None) -> list:
|
||||
"""Return items in minibatch order.
|
||||
|
||||
Uses a deterministic shuffle when a seed is provided so resume runs keep
|
||||
the same minibatch composition. Falls back to input order when no seed is
|
||||
available.
|
||||
"""
|
||||
ordered = list(items)
|
||||
if seed is None:
|
||||
return ordered
|
||||
random.Random(seed).shuffle(ordered)
|
||||
return ordered
|
||||
|
||||
|
||||
def run_minibatch_reflect(
|
||||
results: list[dict],
|
||||
skill_content: str,
|
||||
prediction_dir: str,
|
||||
patches_dir: str,
|
||||
workers: int,
|
||||
failure_only: bool,
|
||||
minibatch_size: int = 8,
|
||||
edit_budget: int = 4,
|
||||
random_seed: int | None = None,
|
||||
*,
|
||||
error_system: str | None = None,
|
||||
success_system: str | None = None,
|
||||
rejection_context: str = "",
|
||||
trajectory_memory_context: str = "",
|
||||
step_buffer_context: str = "",
|
||||
meta_skill_context: str = "",
|
||||
update_mode: str = "patch",
|
||||
skill_aware_reflection: bool | None = None,
|
||||
skill_aware_appendix_source: str | None = None,
|
||||
) -> list[dict | None]:
|
||||
"""Full minibatch reflect stage: group → parallel optimizer calls → patches.
|
||||
|
||||
Separates failure and success trajectories, splits each into minibatches
|
||||
of size M, runs all minibatches in parallel, and saves patch files.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
results : list[dict]
|
||||
Rollout result dicts; see :class:`~skillopt.types.RolloutResult`.
|
||||
skill_content : str
|
||||
Current skill document.
|
||||
prediction_dir : str
|
||||
Path to ``predictions/`` with ``conversation.json`` files.
|
||||
patches_dir : str
|
||||
Path to save per-minibatch patch JSON files.
|
||||
workers : int
|
||||
Max parallel optimizer calls.
|
||||
failure_only : bool
|
||||
If True, skip success trajectories.
|
||||
minibatch_size : int
|
||||
Trajectories per group (M).
|
||||
edit_budget : int
|
||||
Max edits per minibatch (L).
|
||||
random_seed : int | None
|
||||
Optional seed used to shuffle trajectories before minibatch splitting.
|
||||
error_system, success_system : str | None
|
||||
Optional custom prompts. ``None`` = use generic defaults.
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[dict | None]
|
||||
Patch dicts (with ``source_type`` "failure" or "success").
|
||||
"""
|
||||
# Resolve the skill-aware toggle: explicit kwargs win; otherwise fall back
|
||||
# to the process-wide config switch set by the trainer, so the feature is
|
||||
# env-independent and adapters need no per-benchmark wiring.
|
||||
if skill_aware_reflection is None:
|
||||
skill_aware_reflection = is_skill_aware_enabled()
|
||||
if skill_aware_appendix_source is None:
|
||||
skill_aware_appendix_source = get_skill_aware_appendix_source()
|
||||
|
||||
os.makedirs(patches_dir, exist_ok=True)
|
||||
|
||||
# Separate failure / success
|
||||
failures = [r for r in results if not r.get("hard") or float(r.get("hard", 0)) < 1e-9]
|
||||
successes = [r for r in results if r.get("hard")] if not failure_only else []
|
||||
|
||||
failures = _shuffle_for_minibatch(failures, random_seed)
|
||||
successes = _shuffle_for_minibatch(successes, None if random_seed is None else random_seed + 1)
|
||||
|
||||
# Split into minibatches
|
||||
fail_batches = _split_minibatches(failures, minibatch_size)
|
||||
succ_batches = _split_minibatches(successes, minibatch_size)
|
||||
|
||||
n_fail_batches = len(fail_batches)
|
||||
n_succ_batches = len(succ_batches)
|
||||
print(
|
||||
f" [2/6 REFLECT minibatch] "
|
||||
f"failure={len(failures)}→{n_fail_batches} groups "
|
||||
f"success={len(successes)}→{n_succ_batches} groups "
|
||||
f"(M={minibatch_size}, L={edit_budget}, workers={workers})"
|
||||
)
|
||||
|
||||
raw_patches: list[dict | None] = []
|
||||
|
||||
# Resume support: check for already-done minibatch patches
|
||||
pending_fail: list[tuple[int, list[dict]]] = []
|
||||
for idx, batch in enumerate(fail_batches):
|
||||
path = os.path.join(patches_dir, f"minibatch_fail_{idx:03d}.json")
|
||||
if os.path.exists(path):
|
||||
with open(path) as f:
|
||||
raw_patches.append(json.load(f))
|
||||
else:
|
||||
pending_fail.append((idx, batch))
|
||||
|
||||
pending_succ: list[tuple[int, list[dict]]] = []
|
||||
for idx, batch in enumerate(succ_batches):
|
||||
path = os.path.join(patches_dir, f"minibatch_succ_{idx:03d}.json")
|
||||
if os.path.exists(path):
|
||||
with open(path) as f:
|
||||
raw_patches.append(json.load(f))
|
||||
else:
|
||||
pending_succ.append((idx, batch))
|
||||
|
||||
# ── Worker functions ──────────────────────────────────────────────────
|
||||
def _do_fail(idx: int, batch: list[dict]) -> tuple[str, dict | None]:
|
||||
patch = run_error_analyst_minibatch(
|
||||
skill_content, batch, prediction_dir,
|
||||
edit_budget=edit_budget,
|
||||
system_prompt=error_system,
|
||||
step_buffer_context=step_buffer_context,
|
||||
# backward compat fallback
|
||||
rejection_context=rejection_context,
|
||||
trajectory_memory_context=trajectory_memory_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=update_mode,
|
||||
skill_aware_reflection=skill_aware_reflection,
|
||||
)
|
||||
return f"minibatch_fail_{idx:03d}", patch
|
||||
|
||||
def _do_succ(idx: int, batch: list[dict]) -> tuple[str, dict | None]:
|
||||
patch = run_success_analyst_minibatch(
|
||||
skill_content, batch, prediction_dir,
|
||||
edit_budget=edit_budget,
|
||||
system_prompt=success_system,
|
||||
step_buffer_context=step_buffer_context,
|
||||
trajectory_memory_context=trajectory_memory_context,
|
||||
meta_skill_context=meta_skill_context,
|
||||
update_mode=update_mode,
|
||||
skill_aware_reflection=skill_aware_reflection,
|
||||
emit_appendix_notes=(skill_aware_appendix_source != "failure_only"),
|
||||
)
|
||||
return f"minibatch_succ_{idx:03d}", patch
|
||||
|
||||
# Run all pending minibatches in parallel
|
||||
all_pending = (
|
||||
[("fail", idx, batch) for idx, batch in pending_fail]
|
||||
+ [("succ", idx, batch) for idx, batch in pending_succ]
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=workers) as ex:
|
||||
futs = {}
|
||||
for kind, idx, batch in all_pending:
|
||||
if kind == "fail":
|
||||
futs[ex.submit(_do_fail, idx, batch)] = (kind, idx, len(batch))
|
||||
else:
|
||||
futs[ex.submit(_do_succ, idx, batch)] = (kind, idx, len(batch))
|
||||
|
||||
for i, fut in enumerate(as_completed(futs), 1):
|
||||
kind, idx, batch_len = futs[fut]
|
||||
tag, patch = fut.result()
|
||||
if patch:
|
||||
path = os.path.join(patches_dir, f"{tag}.json")
|
||||
with open(path, "w") as f:
|
||||
json.dump(patch, f, ensure_ascii=False, indent=2)
|
||||
raw_patches.append(patch)
|
||||
n_edits = len(get_payload_items(patch.get("patch", {}) if patch else {}, update_mode))
|
||||
print(
|
||||
f" [analyst] {i}/{len(all_pending)} {tag} "
|
||||
f"({batch_len} trajs) → {n_edits} {payload_label(update_mode)}"
|
||||
)
|
||||
|
||||
return raw_patches
|
||||
Reference in New Issue
Block a user