233 lines
8.1 KiB
Python
233 lines
8.1 KiB
Python
"""Tests for PromptOptimizer - BootstrapPromptOptimizer, LLMPromptOptimizer, factory"""
|
|
|
|
import pytest
|
|
|
|
from agentkit.evolution.prompt_optimizer import (
|
|
BootstrapPromptOptimizer,
|
|
LLMPromptOptimizer,
|
|
Module,
|
|
PromptOptimizer,
|
|
Signature,
|
|
create_prompt_optimizer,
|
|
)
|
|
|
|
|
|
def _make_module(instruction: str = "Find the best result.") -> Module:
|
|
return Module(
|
|
name="test_module",
|
|
signature=Signature(
|
|
input_fields={"query": "search query"},
|
|
output_fields={"result": "search result"},
|
|
instruction=instruction,
|
|
),
|
|
)
|
|
|
|
|
|
# ── BootstrapPromptOptimizer ───────────────────────────────────
|
|
|
|
|
|
class TestBootstrapPromptOptimizer:
|
|
"""测试 BootstrapPromptOptimizer"""
|
|
|
|
def test_is_alias_for_prompt_optimizer(self):
|
|
"""PromptOptimizer 是 BootstrapPromptOptimizer 的别名"""
|
|
assert PromptOptimizer is BootstrapPromptOptimizer
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_not_enough_examples_returns_unchanged(self):
|
|
"""样本不足时返回未修改的模块"""
|
|
optimizer = BootstrapPromptOptimizer(min_examples_for_optimization=3)
|
|
optimizer.add_example({"q": "1"}, {"a": "1"}, 0.9)
|
|
|
|
result = await optimizer.optimize(_make_module())
|
|
assert result.name == "test_module" # Unchanged
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enough_examples_produces_optimized_module(self):
|
|
"""足够样本时产生优化模块"""
|
|
optimizer = BootstrapPromptOptimizer(max_demos=3, min_examples_for_optimization=2)
|
|
for i in range(3):
|
|
optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)
|
|
|
|
result = await optimizer.optimize(_make_module())
|
|
assert result.name == "test_module_optimized"
|
|
assert len(result.demos) == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failure_examples_add_avoid_patterns(self):
|
|
"""失败样本添加避免模式到指令中"""
|
|
optimizer = BootstrapPromptOptimizer(min_examples_for_optimization=1)
|
|
optimizer.add_example({"q": "good"}, {"a": "good"}, 0.9)
|
|
optimizer.add_example({"bad_input": "bad"}, {"a": "bad"}, 0.3)
|
|
|
|
result = await optimizer.optimize(_make_module())
|
|
assert "Avoid these patterns" in result.signature.instruction
|
|
|
|
def test_example_count(self):
|
|
"""example_count 返回正确的成功/失败数"""
|
|
optimizer = BootstrapPromptOptimizer()
|
|
optimizer.add_example({"q": "1"}, {"a": "1"}, 0.9)
|
|
optimizer.add_example({"q": "2"}, {"a": "2"}, 0.3)
|
|
optimizer.add_example({"q": "3"}, {"a": "3"}, 0.8)
|
|
|
|
success, failure = optimizer.example_count
|
|
assert success == 2
|
|
assert failure == 1
|
|
|
|
|
|
# ── LLMPromptOptimizer ─────────────────────────────────────────
|
|
|
|
|
|
class MockLLMResponse:
|
|
"""Mock LLM response"""
|
|
def __init__(self, content: str):
|
|
self.content = content
|
|
|
|
|
|
class MockLLMGateway:
|
|
"""Mock LLM Gateway"""
|
|
def __init__(self, response_content: str = "Improved instruction for better results."):
|
|
self._response = response_content
|
|
self.chat_called = False
|
|
|
|
async def chat(self, messages, model="default", agent_name="", task_type=""):
|
|
self.chat_called = True
|
|
return MockLLMResponse(self._response)
|
|
|
|
|
|
class FailingLLMGateway:
|
|
"""LLM Gateway that always fails"""
|
|
async def chat(self, messages, **kwargs):
|
|
raise RuntimeError("LLM unavailable")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_optimizer_generates_improved_instruction():
|
|
"""LLMPromptOptimizer 生成改进的指令"""
|
|
gateway = MockLLMGateway()
|
|
optimizer = LLMPromptOptimizer(llm_gateway=gateway)
|
|
|
|
# Add enough examples for bootstrap post-processing
|
|
for i in range(3):
|
|
optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)
|
|
|
|
module = _make_module()
|
|
result = await optimizer.optimize(module)
|
|
|
|
assert result.name == "test_module_optimized"
|
|
assert result.signature.instruction == "Improved instruction for better results."
|
|
assert gateway.chat_called is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_optimizer_falls_back_to_bootstrap_on_failure():
|
|
"""LLM 调用失败时回退到 BootstrapPromptOptimizer"""
|
|
gateway = FailingLLMGateway()
|
|
optimizer = LLMPromptOptimizer(llm_gateway=gateway)
|
|
|
|
# Add enough examples for bootstrap fallback
|
|
for i in range(3):
|
|
optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)
|
|
|
|
module = _make_module()
|
|
result = await optimizer.optimize(module)
|
|
|
|
# Should fall back to bootstrap optimization
|
|
assert result.name == "test_module_optimized"
|
|
assert len(result.demos) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_optimizer_with_reflection_context():
|
|
"""LLMPromptOptimizer 传递反思上下文"""
|
|
from agentkit.evolution.reflector import Reflection
|
|
|
|
gateway = MockLLMGateway()
|
|
optimizer = LLMPromptOptimizer(llm_gateway=gateway)
|
|
|
|
for i in range(3):
|
|
optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)
|
|
|
|
reflection = Reflection(
|
|
task_id="test-001",
|
|
agent_name="test_agent",
|
|
outcome="failure",
|
|
quality_score=0.3,
|
|
patterns=["slow_execution"],
|
|
insights=["Low quality score"],
|
|
suggestions=["Optimize prompt"],
|
|
)
|
|
|
|
module = _make_module()
|
|
result = await optimizer.optimize(module, trace=None, reflection=reflection)
|
|
|
|
assert result.name == "test_module_optimized"
|
|
assert gateway.chat_called is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_optimizer_empty_response_falls_back():
|
|
"""LLM 返回空响应时回退到 bootstrap"""
|
|
gateway = MockLLMGateway(response_content=" ")
|
|
optimizer = LLMPromptOptimizer(llm_gateway=gateway)
|
|
|
|
for i in range(3):
|
|
optimizer.add_example({"q": f"q_{i}"}, {"a": f"a_{i}"}, 0.9)
|
|
|
|
module = _make_module()
|
|
result = await optimizer.optimize(module)
|
|
|
|
# Should fall back to bootstrap
|
|
assert result.name == "test_module_optimized"
|
|
|
|
|
|
def test_llm_optimizer_example_count():
|
|
"""LLMPromptOptimizer 的 example_count 委托给 bootstrap"""
|
|
optimizer = LLMPromptOptimizer(llm_gateway=MockLLMGateway())
|
|
optimizer.add_example({"q": "1"}, {"a": "1"}, 0.9)
|
|
optimizer.add_example({"q": "2"}, {"a": "2"}, 0.3)
|
|
|
|
success, failure = optimizer.example_count
|
|
assert success == 1
|
|
assert failure == 1
|
|
|
|
|
|
# ── Factory function ───────────────────────────────────────────
|
|
|
|
|
|
class TestCreatePromptOptimizer:
|
|
"""测试 create_prompt_optimizer 工厂函数"""
|
|
|
|
def test_bootstrap_type(self):
|
|
"""bootstrap 类型返回 BootstrapPromptOptimizer"""
|
|
optimizer = create_prompt_optimizer("bootstrap")
|
|
assert isinstance(optimizer, BootstrapPromptOptimizer)
|
|
|
|
def test_llm_type_with_gateway(self):
|
|
"""llm 类型有 gateway 时返回 LLMPromptOptimizer"""
|
|
gateway = MockLLMGateway()
|
|
optimizer = create_prompt_optimizer("llm", llm_gateway=gateway)
|
|
assert isinstance(optimizer, LLMPromptOptimizer)
|
|
|
|
def test_llm_type_without_gateway_falls_back(self):
|
|
"""llm 类型无 gateway 时回退到 BootstrapPromptOptimizer"""
|
|
optimizer = create_prompt_optimizer("llm", llm_gateway=None)
|
|
assert isinstance(optimizer, BootstrapPromptOptimizer)
|
|
|
|
def test_auto_type_with_gateway(self):
|
|
"""auto 类型有 gateway 时返回 LLMPromptOptimizer"""
|
|
gateway = MockLLMGateway()
|
|
optimizer = create_prompt_optimizer("auto", llm_gateway=gateway)
|
|
assert isinstance(optimizer, LLMPromptOptimizer)
|
|
|
|
def test_auto_type_without_gateway(self):
|
|
"""auto 类型无 gateway 时返回 BootstrapPromptOptimizer"""
|
|
optimizer = create_prompt_optimizer("auto", llm_gateway=None)
|
|
assert isinstance(optimizer, BootstrapPromptOptimizer)
|
|
|
|
def test_kwargs_passed_through(self):
|
|
"""额外参数传递给优化器"""
|
|
optimizer = create_prompt_optimizer("bootstrap", max_demos=3, min_examples_for_optimization=2)
|
|
assert optimizer._max_demos == 3
|
|
assert optimizer._min_examples == 2
|