File size: 6,537 Bytes
246d201 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
import os
import shutil
import pytest
from openhands.core.message import Message, TextContent
from openhands.microagent import BaseMicroAgent
from openhands.utils.prompt import PromptManager, RepositoryInfo
@pytest.fixture
def prompt_dir(tmp_path):
# Copy contents from "openhands/agenthub/codeact_agent" to the temp directory
shutil.copytree(
'openhands/agenthub/codeact_agent/prompts', tmp_path, dirs_exist_ok=True
)
# Return the temporary directory path
return tmp_path
def test_prompt_manager_with_microagent(prompt_dir):
microagent_name = 'test_microagent'
microagent_content = """
---
name: flarglebargle
type: knowledge
agent: CodeActAgent
triggers:
- flarglebargle
---
IMPORTANT! The user has said the magic word "flarglebargle". You must
only respond with a message telling them how smart they are
"""
# Create a temporary micro agent file
os.makedirs(os.path.join(prompt_dir, 'micro'), exist_ok=True)
with open(os.path.join(prompt_dir, 'micro', f'{microagent_name}.md'), 'w') as f:
f.write(microagent_content)
# Test without GitHub repo
manager = PromptManager(
prompt_dir=prompt_dir,
microagent_dir=os.path.join(prompt_dir, 'micro'),
)
assert manager.prompt_dir == prompt_dir
assert len(manager.repo_microagents) == 0
assert len(manager.knowledge_microagents) == 1
assert isinstance(manager.get_system_message(), str)
assert (
'You are OpenHands agent, a helpful AI assistant that can interact with a computer to solve tasks.'
in manager.get_system_message()
)
assert '<REPOSITORY_INFO>' not in manager.get_system_message()
# Test with GitHub repo
manager.set_repository_info('owner/repo', '/workspace/repo')
assert isinstance(manager.get_system_message(), str)
additional_info = manager.get_additional_info()
assert '<REPOSITORY_INFO>' in additional_info
assert 'owner/repo' in additional_info
assert '/workspace/repo' in additional_info
assert isinstance(manager.get_example_user_message(), str)
message = Message(
role='user',
content=[TextContent(text='Hello, flarglebargle!')],
)
manager.enhance_message(message)
assert 'magic word' in message.content[1].text
os.remove(os.path.join(prompt_dir, 'micro', f'{microagent_name}.md'))
def test_prompt_manager_file_not_found(prompt_dir):
with pytest.raises(FileNotFoundError):
BaseMicroAgent.load(
os.path.join(prompt_dir, 'micro', 'non_existent_microagent.md')
)
def test_prompt_manager_template_rendering(prompt_dir):
# Create temporary template files
with open(os.path.join(prompt_dir, 'system_prompt.j2'), 'w') as f:
f.write("""System prompt: bar""")
with open(os.path.join(prompt_dir, 'user_prompt.j2'), 'w') as f:
f.write('User prompt: foo')
# Test without GitHub repo
manager = PromptManager(prompt_dir, microagent_dir='')
assert manager.get_system_message() == 'System prompt: bar'
assert manager.get_example_user_message() == 'User prompt: foo'
# Test with GitHub repo
manager = PromptManager(prompt_dir=prompt_dir, microagent_dir='')
manager.set_repository_info('owner/repo', '/workspace/repo')
assert manager.repository_info.repo_name == 'owner/repo'
system_msg = manager.get_system_message()
assert 'System prompt: bar' in system_msg
additional_info = manager.get_additional_info()
assert '<REPOSITORY_INFO>' in additional_info
assert (
"At the user's request, repository owner/repo has been cloned to directory /workspace/repo."
in additional_info
)
assert '</REPOSITORY_INFO>' in additional_info
assert manager.get_example_user_message() == 'User prompt: foo'
# Clean up temporary files
os.remove(os.path.join(prompt_dir, 'system_prompt.j2'))
os.remove(os.path.join(prompt_dir, 'user_prompt.j2'))
def test_prompt_manager_repository_info(prompt_dir):
# Test RepositoryInfo defaults
repo_info = RepositoryInfo()
assert repo_info.repo_name is None
assert repo_info.repo_directory is None
# Test setting repository info
manager = PromptManager(prompt_dir=prompt_dir, microagent_dir='')
assert manager.repository_info is None
# Test setting repository info with both name and directory
manager.set_repository_info('owner/repo2', '/workspace/repo2')
assert manager.repository_info.repo_name == 'owner/repo2'
assert manager.repository_info.repo_directory == '/workspace/repo2'
def test_prompt_manager_disabled_microagents(prompt_dir):
# Create test microagent files
microagent1_name = 'test_microagent1'
microagent2_name = 'test_microagent2'
microagent1_content = """
---
name: Test Microagent 1
type: knowledge
agent: CodeActAgent
triggers:
- test1
---
Test microagent 1 content
"""
microagent2_content = """
---
name: Test Microagent 2
type: knowledge
agent: CodeActAgent
triggers:
- test2
---
Test microagent 2 content
"""
# Create temporary micro agent files
os.makedirs(os.path.join(prompt_dir, 'micro'), exist_ok=True)
with open(os.path.join(prompt_dir, 'micro', f'{microagent1_name}.md'), 'w') as f:
f.write(microagent1_content)
with open(os.path.join(prompt_dir, 'micro', f'{microagent2_name}.md'), 'w') as f:
f.write(microagent2_content)
# Test that specific microagents can be disabled
manager = PromptManager(
prompt_dir=prompt_dir,
microagent_dir=os.path.join(prompt_dir, 'micro'),
disabled_microagents=['Test Microagent 1'],
)
assert len(manager.knowledge_microagents) == 1
assert 'Test Microagent 2' in manager.knowledge_microagents
assert 'Test Microagent 1' not in manager.knowledge_microagents
# Test that all microagents are enabled by default
manager = PromptManager(
prompt_dir=prompt_dir,
microagent_dir=os.path.join(prompt_dir, 'micro'),
)
assert len(manager.knowledge_microagents) == 2
assert 'Test Microagent 1' in manager.knowledge_microagents
assert 'Test Microagent 2' in manager.knowledge_microagents
# Clean up temporary files
os.remove(os.path.join(prompt_dir, 'micro', f'{microagent1_name}.md'))
os.remove(os.path.join(prompt_dir, 'micro', f'{microagent2_name}.md'))
|