Spaces:
Build error
Build error
File size: 2,692 Bytes
01523b5 |
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 |
from __future__ import annotations
from typing import TYPE_CHECKING, List
from agentverse.message import Message
from . import selector_registry as SelectorRegistry
from .base import BaseSelector
import json
import re
if TYPE_CHECKING:
from agentverse.environments import BaseEnvironment
def extract(content: str, keyword: str):
result = ""
flag = False
for line in content.split('\n'):
if line.strip().startswith(keyword):
flag = True
continue
if flag:
result += line
result += "\n"
return result
@SelectorRegistry.register("sde_team")
class SdeTeamSelector(BaseSelector):
def select_message(self, environment: BaseEnvironment, messages: List[Message]) -> List[Message]:
last_sender = environment.last_messages[0].sender
selected = messages
if last_sender == "unit_test_generator":
unit_tests = set()
for message in selected:
unit_test = extract(message.content, "<unit test>:")
if unit_test not in unit_tests:
unit_tests.add(extract(message.content, "<unit test>:"))
unit_tests = list(unit_tests)
environment.rule_params["unit_tests"] = str(unit_tests)
new_message = Message(
content="",
sender="unit_test_generator",
receiver=[],
) # TODO: set the content of the message
selected = [new_message]
elif last_sender == "code_writer":
cur_code = extract(selected[0].content, "<code>:")
environment.rule_params["code"] = cur_code
from .code_api import execute_unit_tests
feedback = execute_unit_tests(environment.rule_params["code"], eval(environment.rule_params["unit_tests"]))
environment.rule_params["feedback"] = feedback
selected[0].content = f"<current code>:\n\n{cur_code}\n\n<unit test feedback>:\n{feedback}"
f_dict = json.loads(feedback)
if f_dict["is_passing"]:
environment.rule_params["end_flag"] = True
elif last_sender == "code_reviewer":
code_review = selected[0].content
cur_code = environment.rule_params["code"]
selected[0].content = f"<current code>:\n\n{cur_code}\n\n{code_review}"
feedback = environment.rule_params["feedback"]
f_dict = json.loads(feedback)
if f_dict["is_passing"]:
environment.rule_params["end_flag"] = True
return selected |