298 lines
9.2 KiB
Python
298 lines
9.2 KiB
Python
"""Tests for PromptBuilder."""
|
|
|
|
from alfred.agent.prompts import PromptBuilder
|
|
from alfred.agent.registry import make_tools
|
|
|
|
|
|
class TestPromptBuilder:
|
|
"""Tests for PromptBuilder."""
|
|
|
|
def test_init(self, memory):
|
|
"""Should initialize with tools."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
assert builder.tools is tools
|
|
|
|
def test_build_system_prompt(self, memory):
|
|
"""Should build a complete system prompt."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "AI assistant" in prompt
|
|
assert "media library" in prompt
|
|
assert "AVAILABLE TOOLS" in prompt
|
|
|
|
def test_includes_tools(self, memory):
|
|
"""Should include all tool descriptions."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
for tool_name in tools.keys():
|
|
assert tool_name in prompt
|
|
|
|
def test_includes_config(self, memory):
|
|
"""Should include current configuration."""
|
|
memory.ltm.set_config("download_folder", "/path/to/downloads")
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "/path/to/downloads" in prompt
|
|
|
|
def test_includes_search_results(self, memory_with_search_results):
|
|
"""Should include search results summary."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "LAST SEARCH" in prompt
|
|
assert "Inception 1080p" in prompt
|
|
assert "3 results" in prompt or "results available" in prompt
|
|
|
|
def test_includes_search_result_names(self, memory_with_search_results):
|
|
"""Should include search result names."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "Inception.2010.1080p.BluRay.x264" in prompt
|
|
|
|
def test_includes_active_downloads(self, memory):
|
|
"""Should include active downloads."""
|
|
memory.episodic.add_active_download(
|
|
{
|
|
"task_id": "123",
|
|
"name": "Test.Movie.mkv",
|
|
"progress": 50,
|
|
}
|
|
)
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "ACTIVE DOWNLOADS" in prompt
|
|
assert "Test.Movie.mkv" in prompt
|
|
|
|
def test_includes_pending_question(self, memory):
|
|
"""Should include pending question."""
|
|
memory.episodic.set_pending_question(
|
|
"Which torrent?",
|
|
[{"index": 1, "label": "Option 1"}, {"index": 2, "label": "Option 2"}],
|
|
{},
|
|
)
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "PENDING QUESTION" in prompt
|
|
assert "Which torrent?" in prompt
|
|
|
|
def test_includes_last_error(self, memory):
|
|
"""Should include last error."""
|
|
memory.episodic.add_error("find_torrent", "API timeout")
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "RECENT ERRORS" in prompt
|
|
assert "API timeout" in prompt
|
|
|
|
def test_includes_workflow(self, memory):
|
|
"""Should include current workflow."""
|
|
memory.stm.start_workflow("download", {"title": "Inception"})
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "CURRENT WORKFLOW" in prompt
|
|
assert "download" in prompt
|
|
|
|
def test_includes_topic(self, memory):
|
|
"""Should include current topic."""
|
|
memory.stm.set_topic("selecting_torrent")
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "CURRENT TOPIC" in prompt
|
|
assert "selecting_torrent" in prompt
|
|
|
|
def test_includes_entities(self, memory):
|
|
"""Should include extracted entities."""
|
|
memory.stm.set_entity("movie_title", "Inception")
|
|
memory.stm.set_entity("year", 2010)
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "EXTRACTED ENTITIES" in prompt
|
|
assert "Inception" in prompt
|
|
|
|
def test_includes_rules(self, memory):
|
|
"""Should include important rules."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "IMPORTANT RULES" in prompt
|
|
assert "add_torrent_by_index" in prompt
|
|
|
|
def test_includes_examples(self, memory):
|
|
"""Should include usage examples."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
assert "EXAMPLES" in prompt
|
|
assert "download the 3rd one" in prompt or "torrent number" in prompt
|
|
|
|
def test_empty_context(self, memory):
|
|
"""Should handle empty context gracefully."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
# Should not crash and should have basic structure
|
|
assert "AVAILABLE TOOLS" in prompt
|
|
assert "CURRENT CONFIGURATION" in prompt
|
|
|
|
def test_limits_search_results_display(self, memory):
|
|
"""Should limit displayed search results."""
|
|
# Add many results
|
|
results = [{"name": f"Torrent {i}", "seeders": i} for i in range(20)]
|
|
memory.episodic.store_search_results("test", results)
|
|
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
prompt = builder.build_system_prompt()
|
|
|
|
# Should show first 5 and indicate more
|
|
assert "Torrent 0" in prompt or "1." in prompt
|
|
assert "... and" in prompt or "more" in prompt
|
|
|
|
# REMOVED: test_json_format_in_prompt
|
|
# We removed the "action" format from prompts as it was confusing the LLM
|
|
# The LLM now uses native OpenAI tool calling format
|
|
|
|
|
|
class TestFormatToolsDescription:
|
|
"""Tests for _format_tools_description method."""
|
|
|
|
def test_format_all_tools(self, memory):
|
|
"""Should format all tools."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
desc = builder._format_tools_description()
|
|
|
|
for tool in tools.values():
|
|
assert tool.name in desc
|
|
assert tool.description in desc
|
|
|
|
def test_includes_parameters(self, memory):
|
|
"""Should include parameter schemas."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
desc = builder._format_tools_description()
|
|
|
|
assert "Parameters:" in desc
|
|
assert '"type"' in desc
|
|
|
|
|
|
class TestFormatEpisodicContext:
|
|
"""Tests for _format_episodic_context method."""
|
|
|
|
def test_empty_episodic(self, memory):
|
|
"""Should return empty string for empty episodic."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
context = builder._format_episodic_context(memory)
|
|
|
|
assert context == ""
|
|
|
|
def test_with_search_results(self, memory_with_search_results):
|
|
"""Should format search results."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
context = builder._format_episodic_context(memory_with_search_results)
|
|
|
|
assert "LAST SEARCH" in context
|
|
assert "Inception 1080p" in context
|
|
|
|
def test_with_multiple_sections(self, memory):
|
|
"""Should format multiple sections."""
|
|
memory.episodic.store_search_results("test", [{"name": "Result"}])
|
|
memory.episodic.add_active_download({"task_id": "1", "name": "Download"})
|
|
memory.episodic.add_error("action", "error")
|
|
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
context = builder._format_episodic_context(memory)
|
|
|
|
assert "LAST SEARCH" in context
|
|
assert "ACTIVE DOWNLOADS" in context
|
|
assert "RECENT ERRORS" in context
|
|
|
|
|
|
class TestFormatStmContext:
|
|
"""Tests for _format_stm_context method."""
|
|
|
|
def test_empty_stm(self, memory):
|
|
"""Should return language info even for empty STM."""
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
context = builder._format_stm_context(memory)
|
|
|
|
# Should at least show language
|
|
assert "CONVERSATION LANGUAGE" in context or context == ""
|
|
|
|
def test_with_workflow(self, memory):
|
|
"""Should format workflow."""
|
|
memory.stm.start_workflow("download", {"title": "Test"})
|
|
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
context = builder._format_stm_context(memory)
|
|
|
|
assert "CURRENT WORKFLOW" in context
|
|
assert "download" in context
|
|
|
|
def test_with_all_sections(self, memory):
|
|
"""Should format all STM sections."""
|
|
memory.stm.start_workflow("download", {"title": "Test"})
|
|
memory.stm.set_topic("searching")
|
|
memory.stm.set_entity("key", "value")
|
|
|
|
tools = make_tools()
|
|
builder = PromptBuilder(tools)
|
|
|
|
context = builder._format_stm_context(memory)
|
|
|
|
assert "CURRENT WORKFLOW" in context
|
|
assert "CURRENT TOPIC" in context
|
|
assert "EXTRACTED ENTITIES" in context
|