fix extract_content_from_output

This commit is contained in:
better629 2023-11-23 11:21:25 +08:00
parent 1a67878b7e
commit 502bb2c449
2 changed files with 9 additions and 4 deletions

View file

@ -275,13 +275,15 @@ def extract_content_from_output(content: str, right_key: str = "[/CONTENT]"):
# TODO find a more general pattern
# # for `[CONTENT]xxx[CONTENT]xxxx[/CONTENT] situation
logger.warning(f"extract_content try another pattern: {pattern}")
raw_content = copy.deepcopy(new_content + right_key)
if right_key not in new_content:
raw_content = copy.deepcopy(new_content + "\n" + right_key)
# # pattern = r"\[CONTENT\](\s*\{.*?\}\s*)\[/CONTENT\]"
new_content = re_extract_content(raw_content, pattern)
else:
if right_key in new_content:
idx = new_content.find(right_key)
new_content = new_content[:idx]
new_content = new_content.strip()
return new_content

View file

@ -216,7 +216,8 @@ def test_extract_content_from_output():
'UNCLEAR:\n\n* How to start the game.\n\nGreat job! This JSON output should provide a clear and ' \
'comprehensive overview of the project\'s requirements and dependencies.'
output = extract_content_from_output(output)
assert output.startswith('{\n"Required Python third-party packages')
assert output.startswith('{\n"Required Python third-party packages') and \
output.endswith('UNCLEAR": "How to start the game."\n]')
output = 'Sure, I would be happy to help! Here is the information you provided, formatted as a JSON object ' \
'inside the [CONTENT] tag:\n\n[CONTENT]\n{\n"Original Requirements": "Create a 2048 game",\n"Search ' \
@ -245,7 +246,8 @@ def test_extract_content_from_output():
'[/CONTENT]\n\nI hope this helps! Let me know if you have any further questions or if there anything ' \
'else I can do to assist you.'
output = extract_content_from_output(output)
assert output.startswith('{\n"Original Requirements"')
assert output.startswith('{\n"Original Requirements"') and \
output.endswith('"Anything UNCLEAR": ""\n}')
output = """ Sure, I'd be happy to help! Here's the JSON output for the given context:\n\n[CONTENT]\n{
"Implementation approach": "We will use the open-source framework PyGame to create a 2D game engine, which will
@ -270,4 +272,5 @@ def test_extract_content_from_output():
information for a developer to understand the design and implementation of the 2048 game.
"""
output = extract_content_from_output(output)
assert output.startswith('{\n"Implementation approach"') and "[/CONTENT]" not in output
assert output.startswith('{\n"Implementation approach"') and \
output.endswith('"Anything UNCLEAR": "The requirement is clear to me."\n}')