|
|
from random import randint |
|
|
from typing import Iterator |
|
|
|
|
|
from injector import Injector |
|
|
|
|
|
from taskweaver.logging import LoggingModule |
|
|
from taskweaver.memory import Attachment, Post |
|
|
from taskweaver.memory.attachment import AttachmentType |
|
|
from taskweaver.role import PostTranslator |
|
|
|
|
|
response_str1 = ( |
|
|
'{"response": [{"type": "thought", "content": "This is the thought"}, {"type": "python", ' |
|
|
'"content": "print(\'This is the code\')"}, {"type": "text", "content": "This ' |
|
|
'is the text"}, {"type": "sample", "content": "print(\'This is the ' |
|
|
'sample code\')"}, {"type": "execution_status", "content": "SUCCESS"}, ' |
|
|
'{"type": "execution_result", "content": "This is the execution result"}, ' |
|
|
'{"type": "send_to", "content": "Planner"}, {"type": "message", "content": ' |
|
|
'"This is the message"}]}' |
|
|
) |
|
|
|
|
|
role_name = "ProgramApe" |
|
|
executor_name = "CodeExecutor" |
|
|
|
|
|
app_injector = Injector( |
|
|
[LoggingModule], |
|
|
) |
|
|
translator = app_injector.create_object(PostTranslator) |
|
|
|
|
|
|
|
|
def test_parse_llm_stream(): |
|
|
def response_str() -> Iterator[str]: |
|
|
words = response_str1.split(" ") |
|
|
|
|
|
pos = 0 |
|
|
|
|
|
while True: |
|
|
n = randint(1, 10) |
|
|
part = " ".join(words[pos : pos + n]) + " " |
|
|
yield part |
|
|
pos += n |
|
|
if pos >= len(words): |
|
|
break |
|
|
|
|
|
attachments = translator.parse_llm_output_stream(response_str()) |
|
|
attachment_list = list(attachments) |
|
|
assert len(attachment_list) == 8 |
|
|
|
|
|
|
|
|
def test_parse_llm(): |
|
|
def early_stop(type: AttachmentType, text: str) -> bool: |
|
|
if type in [AttachmentType.python, AttachmentType.sample, AttachmentType.text]: |
|
|
return True |
|
|
return False |
|
|
|
|
|
response = translator.raw_text_to_post( |
|
|
llm_output=response_str1, |
|
|
send_from="CodeInterpreter", |
|
|
event_handler=lambda t, v: print(f"{t}: {v}"), |
|
|
early_stop=early_stop, |
|
|
) |
|
|
|
|
|
assert response.message is None |
|
|
assert response.send_to is None |
|
|
assert response.send_from == "CodeInterpreter" |
|
|
assert len(response.attachment_list) == 2 |
|
|
assert response.attachment_list[0].type == AttachmentType.thought |
|
|
assert response.attachment_list[0].content == "This is the thought" |
|
|
|
|
|
assert response.attachment_list[1].type == AttachmentType.python |
|
|
assert response.attachment_list[1].content == "print('This is the code')" |
|
|
|
|
|
response = translator.raw_text_to_post( |
|
|
llm_output=response_str1, |
|
|
send_from="CodeInterpreter", |
|
|
event_handler=lambda t, v: print(f"{t}: {v}"), |
|
|
) |
|
|
assert len(response.attachment_list) == 6 |
|
|
assert response.attachment_list[4].type == AttachmentType.execution_status |
|
|
assert response.attachment_list[4].content == "SUCCESS" |
|
|
assert response.attachment_list[5].type == AttachmentType.execution_result |
|
|
assert response.attachment_list[5].content == "This is the execution result" |
|
|
|
|
|
|
|
|
def test_post_to_raw_text(): |
|
|
post = Post.create(message="This is the message", send_from="CodeInterpreter", send_to="Planner") |
|
|
|
|
|
prompt = translator.post_to_raw_text(post=post, if_format_message=True, if_format_send_to=True) |
|
|
assert prompt == ( |
|
|
'{"response": [{"type": "send_to", "content": "Planner"}, {"type": "message", ' |
|
|
'"content": "This is the message"}]}' |
|
|
) |
|
|
|
|
|
prompt = translator.post_to_raw_text(post=post, if_format_message=False, if_format_send_to=False) |
|
|
assert prompt == '{"response": []}' |
|
|
|
|
|
post.add_attachment(Attachment.create(type="thought", content="This is the thought")) |
|
|
post.add_attachment(Attachment.create(type="python", content="print('This is the code')")) |
|
|
post.add_attachment(Attachment.create(type="text", content="This is the text")) |
|
|
post.add_attachment(Attachment.create(type="sample", content="print('This is the sample code')")) |
|
|
post.add_attachment(Attachment.create(type="execution_status", content="SUCCESS")) |
|
|
post.add_attachment(Attachment.create(type="execution_result", content="This is the execution result")) |
|
|
|
|
|
prompt = translator.post_to_raw_text(post=post, if_format_message=True, if_format_send_to=True) |
|
|
assert prompt == response_str1 |
|
|
|
