File size: 4,663 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 |
import json
import os
from unittest.mock import MagicMock
import pytest
import yaml
from pytest import TempPathFactory
from openhands.agenthub.micro.registry import all_microagents
from openhands.controller.agent import Agent
from openhands.controller.state.state import State
from openhands.core.config import AgentConfig
from openhands.events.action import MessageAction
from openhands.events.stream import EventStream
from openhands.storage import get_file_store
@pytest.fixture
def temp_dir(tmp_path_factory: TempPathFactory) -> str:
return str(tmp_path_factory.mktemp('test_micro_agents'))
@pytest.fixture
def event_stream(temp_dir):
file_store = get_file_store('local', temp_dir)
event_stream = EventStream('asdf', file_store)
yield event_stream
@pytest.fixture
def agent_configs():
return {
'CoderAgent': AgentConfig(memory_enabled=True),
'BrowsingAgent': AgentConfig(memory_enabled=True),
}
def test_all_agents_are_loaded():
assert all_microagents is not None
assert len(all_microagents) > 1
base = os.path.join('openhands', 'agenthub', 'micro')
full_path = os.path.dirname(__file__) + '/../../' + base
agent_names = set()
for root, _, files in os.walk(full_path):
for file in files:
if file == 'agent.yaml':
file_path = os.path.join(root, file)
with open(file_path, 'r') as yaml_file:
data = yaml.safe_load(yaml_file)
agent_names.add(data['name'])
assert agent_names == set(all_microagents.keys())
def test_coder_agent_with_summary(event_stream: EventStream, agent_configs: dict):
"""Coder agent should render code summary as part of prompt"""
mock_llm = MagicMock()
content = json.dumps({'action': 'finish', 'args': {}})
mock_llm.completion.return_value = {'choices': [{'message': {'content': content}}]}
mock_llm.format_messages_for_llm.return_value = [
{
'role': 'user',
'content': "This is a dummy task. This is a dummy summary about this repo. Here's a summary of the codebase, as it relates to this task.",
}
]
coder_agent = Agent.get_cls('CoderAgent')(
llm=mock_llm, config=agent_configs['CoderAgent']
)
assert coder_agent is not None
# give it some history
task = 'This is a dummy task'
history = list()
history.append(MessageAction(content=task))
summary = 'This is a dummy summary about this repo'
state = State(history=history, inputs={'summary': summary})
coder_agent.step(state)
mock_llm.completion.assert_called_once()
_, kwargs = mock_llm.completion.call_args
prompt_element = kwargs['messages'][0]['content']
if isinstance(prompt_element, dict):
prompt = prompt_element['content']
else:
prompt = prompt_element
assert task in prompt
assert "Here's a summary of the codebase, as it relates to this task" in prompt
assert summary in prompt
def test_coder_agent_without_summary(event_stream: EventStream, agent_configs: dict):
"""When there's no codebase_summary available, there shouldn't be any prompt
about 'code summary'
"""
mock_llm = MagicMock()
content = json.dumps({'action': 'finish', 'args': {}})
mock_llm.completion.return_value = {'choices': [{'message': {'content': content}}]}
mock_llm.format_messages_for_llm.return_value = [
{
'role': 'user',
'content': [
{
'type': 'text',
'text': "This is a dummy task. This is a dummy summary about this repo. Here's a summary of the codebase, as it relates to this task.",
}
],
}
]
coder_agent = Agent.get_cls('CoderAgent')(
llm=mock_llm, config=agent_configs['CoderAgent']
)
assert coder_agent is not None
# give it some history
task = 'This is a dummy task'
history = list()
history.append(MessageAction(content=task))
# set state without codebase summary
state = State(history=history)
coder_agent.step(state)
mock_llm.completion.assert_called_once()
_, kwargs = mock_llm.completion.call_args
prompt_element = kwargs['messages'][0]['content']
if isinstance(prompt_element, dict):
prompt = prompt_element['content']
else:
prompt = prompt_element
print(f'\n{prompt_element}\n')
assert "Here's a summary of the codebase, as it relates to this task" not in prompt
|