mirror of
https://github.com/kenhuangus/wikiskill
synced 2026-09-22 23:34:25 +08:00
100 lines
4.2 KiB
Python
100 lines
4.2 KiB
Python
from wikiskill.datasets import make_demo_dataset, _demo_tools
|
|
from wikiskill.llm import MockLLM, extract_json
|
|
from wikiskill.skills import SkillsLayer
|
|
from wikiskill.workspace import Workspace
|
|
|
|
|
|
def _ws(tmp_path):
|
|
return Workspace(str(tmp_path / "ws"))
|
|
|
|
|
|
def test_extract_json():
|
|
assert extract_json('```json\n{"a": 1}\n```') == {"a": 1}
|
|
assert extract_json('noise {"a": {"b": 2}} tail') == {"a": {"b": 2}}
|
|
assert extract_json("no json") is None
|
|
|
|
|
|
def test_skills_create_snapshot_restore(tmp_path):
|
|
skills = SkillsLayer(_ws(tmp_path))
|
|
skills.apply_proposal({"action": "create", "skill": "s1",
|
|
"content": "# S1\nline1\nline2\n", "purpose": "why"})
|
|
assert skills.list_skills() == ["s1"]
|
|
snap = skills.snapshot()
|
|
skills.apply_proposal({"action": "patch", "skill": "s1",
|
|
"edits": [{"op": "replace", "target": "line2",
|
|
"content": "LINE2"}], "purpose": "why2"})
|
|
assert "LINE2" in skills.skill_content("s1")
|
|
skills.restore(snap)
|
|
assert "LINE2" not in skills.skill_content("s1")
|
|
skills.apply_proposal({"action": "create", "skill": "s2", "content": "x",
|
|
"purpose": "p"})
|
|
skills.restore(snap)
|
|
assert skills.list_skills() == ["s1"]
|
|
|
|
|
|
def test_skills_restore_reverts_purpose_md(tmp_path):
|
|
"""A rejected patch must roll back BOTH SKILL.md and PURPOSE.md exactly
|
|
to S_{k-1} (skills-only rollback, §3.2.4)."""
|
|
skills = SkillsLayer(_ws(tmp_path))
|
|
skills.apply_proposal({"action": "create", "skill": "s",
|
|
"content": "body", "purpose": "v1"})
|
|
snap = skills.snapshot()
|
|
skills.apply_proposal({"action": "patch", "skill": "s",
|
|
"edits": [{"op": "append", "content": "extra"}],
|
|
"purpose": "v2 (append)"})
|
|
assert "v2" in skills.purpose("s")
|
|
skills.restore(snap)
|
|
assert skills.purpose("s") == "v1\n"
|
|
assert skills.skill_content("s") == "body"
|
|
|
|
|
|
def test_skill_injection_context(tmp_path):
|
|
skills = SkillsLayer(_ws(tmp_path))
|
|
assert "(No active skills.)" in skills.active_skills_context()
|
|
skills.apply_proposal({"action": "create", "skill": "s", "content": "BODY",
|
|
"purpose": "p"})
|
|
assert "BODY" in skills.active_skills_context()
|
|
|
|
|
|
def test_mock_inference_without_and_with_skill(tmp_path):
|
|
from wikiskill.agents import InferenceAgent
|
|
ws = _ws(tmp_path)
|
|
skills = SkillsLayer(ws)
|
|
agent = InferenceAgent(MockLLM(), _demo_tools(), ws, skills)
|
|
ds = make_demo_dataset()
|
|
traj = agent.run(ds.val[0])
|
|
assert traj.correct is False # truncation error without skill
|
|
skills.apply_proposal({"action": "create", "skill": "unit-conversion",
|
|
"content": "Round ONLY the final result to 2 decimal places.",
|
|
"purpose": "p"})
|
|
agent2 = InferenceAgent(MockLLM(), _demo_tools(), ws, skills)
|
|
traj2 = agent2.run(ds.val[0])
|
|
assert traj2.correct is True
|
|
|
|
|
|
def test_stratified_sample():
|
|
from wikiskill.agents import stratified_sample
|
|
from wikiskill.datasets import Trajectory
|
|
|
|
def t(i, ok):
|
|
return Trajectory(task_id=i, messages=[{"role": "user", "content": "x"}],
|
|
tool_calls=[], final_answer="a", correct=ok)
|
|
traces = [t(f"f{i}", False) for i in range(9)] + [t(f"p{i}", True) for i in range(5)]
|
|
s = stratified_sample(traces)
|
|
assert sum(1 for x in s if not x.correct) == 5
|
|
assert sum(1 for x in s if x.correct) == 3
|
|
|
|
|
|
def test_stratified_sample_caps_whole_log():
|
|
"""Paper App. C: the ENTIRE execution log is capped at 15,000 chars, not just
|
|
the last message."""
|
|
from wikiskill.agents import stratified_sample
|
|
from wikiskill.datasets import Trajectory
|
|
big = Trajectory(task_id="big", correct=False, tool_calls=[], final_answer=None,
|
|
prediction=None,
|
|
messages=[{"role": "system", "content": "s" * 9000},
|
|
{"role": "assistant", "content": "a" * 9000}])
|
|
(sampled,) = stratified_sample([big])
|
|
assert sum(len(str(m.get("content", ""))) for m in sampled.messages) < 16000
|
|
assert any("[truncated]" in str(m.get("content", "")) for m in sampled.messages)
|