149 lines
4.7 KiB
Python
149 lines
4.7 KiB
Python
"""Tests for prompt optimization (M9.1)."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
from cmdforge.prompt_optimizer import (
|
|
PromptVariation,
|
|
OptimizationResult,
|
|
generate_variations,
|
|
optimize_tool,
|
|
_get_strategies,
|
|
_truncate,
|
|
)
|
|
|
|
|
|
class TestTruncate:
|
|
def test_short_text_unchanged(self):
|
|
assert _truncate("hello", 10) == "hello"
|
|
|
|
def test_long_text_truncated(self):
|
|
result = _truncate("this is a very long sentence that needs truncation", 20)
|
|
assert len(result) <= 20
|
|
assert result.endswith("...")
|
|
|
|
|
|
class TestGenerateVariations:
|
|
def test_generates_for_prompt_steps(self, tmp_path):
|
|
from cmdforge.tool import Tool, PromptStep
|
|
|
|
with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
|
|
tool = Tool(
|
|
name="test-tool",
|
|
steps=[
|
|
PromptStep(
|
|
prompt="Analyze the input and produce a summary",
|
|
provider="mock",
|
|
output_var="summary",
|
|
),
|
|
],
|
|
output="{summary}",
|
|
)
|
|
variations = generate_variations(tool, count=3)
|
|
assert len(variations) >= 1
|
|
for v in variations:
|
|
assert v.step_index == 0
|
|
assert v.original == "Analyze the input and produce a summary"
|
|
assert v.variation != v.original
|
|
|
|
def test_skips_non_prompt_steps(self, tmp_path):
|
|
from cmdforge.tool import Tool, CodeStep
|
|
|
|
with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
|
|
tool = Tool(
|
|
name="test-tool",
|
|
steps=[CodeStep(code="print('hello')", output_var="out")],
|
|
output="{out}",
|
|
)
|
|
variations = generate_variations(tool)
|
|
assert len(variations) == 0
|
|
|
|
|
|
class TestGetStrategies:
|
|
def test_produces_variations(self):
|
|
strategies = _get_strategies("Summarize the text", 3)
|
|
assert len(strategies) == 3
|
|
for strat, text in strategies:
|
|
assert len(text) > 0
|
|
assert strat in ("focused", "concise", "detailed", "rephrase", "structure")
|
|
|
|
def test_truncates_concise(self):
|
|
original = "A" * 200
|
|
strategies = _get_strategies(original, 3)
|
|
concise = [t for s, t in strategies if s == "concise"]
|
|
if concise:
|
|
assert len(concise[0]) <= 125 # 120 + "..."
|
|
|
|
|
|
class TestOptimizeTool:
|
|
def test_baseline_score_recorded(self, tmp_path):
|
|
from cmdforge.tool import Tool, PromptStep
|
|
|
|
with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
|
|
tool = Tool(
|
|
name="opt-test",
|
|
steps=[
|
|
PromptStep(
|
|
prompt="Summarize: {input}",
|
|
provider="mock",
|
|
output_var="result",
|
|
),
|
|
],
|
|
output="{result}",
|
|
)
|
|
result = optimize_tool(tool, count=2)
|
|
assert result.tool_name == "opt-test"
|
|
assert result.baseline_score >= 0
|
|
assert len(result.variations) >= 1
|
|
assert len(result.scores) == len(result.variations)
|
|
|
|
def test_no_prompt_steps_returns_empty(self, tmp_path):
|
|
from cmdforge.tool import Tool, CodeStep
|
|
|
|
with patch("cmdforge.tool.TOOLS_DIR", tmp_path / ".cmdforge"):
|
|
tool = Tool(
|
|
name="code-only",
|
|
steps=[CodeStep(code="pass", output_var="x")],
|
|
output="{x}",
|
|
)
|
|
result = optimize_tool(tool)
|
|
assert len(result.variations) == 0
|
|
assert result.best is None
|
|
|
|
|
|
class TestPromptVariation:
|
|
def test_to_dict(self):
|
|
v = PromptVariation(
|
|
step_index=0,
|
|
original="hello",
|
|
variation="hi there",
|
|
strategy="rephrase",
|
|
)
|
|
d = v.to_dict()
|
|
assert d["original"] == "hello"
|
|
assert d["strategy"] == "rephrase"
|
|
|
|
|
|
class TestOptimizationResult:
|
|
def test_improvement_calculation(self):
|
|
result = OptimizationResult(
|
|
tool_name="test",
|
|
variations=[
|
|
PromptVariation(0, "a", "b", "rephrase"),
|
|
],
|
|
scores={0: 5},
|
|
best_index=0,
|
|
baseline_score=3,
|
|
)
|
|
assert result.best is not None
|
|
assert result.improvement > 0
|
|
|
|
def test_no_improvement_when_no_baseline(self):
|
|
result = OptimizationResult(
|
|
tool_name="test",
|
|
variations=[],
|
|
scores={},
|
|
baseline_score=0,
|
|
)
|
|
assert result.best is None
|
|
assert result.improvement == 0.0
|