Spaces:
Sleeping
Sleeping
File size: 3,763 Bytes
acd245a |
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 |
"""Module to parse the output of the Aider model."""
import re
from dataclasses import dataclass
from typing import List
@dataclass
class AiderOutputBlock:
"""A single block of output from the Aider model."""
content: str
type: str # 'text', 'code', 'prompt', 'separator', 'token_info'
is_code_block: bool = False
is_prompt: bool = False
is_separator: bool = False
is_token_info: bool = False
class AiderOutputParser:
"""Output parser for Aider model."""
def __init__(self) -> None:
"""Initialize the output parser."""
self.code_block_pattern = re.compile(
r"<<<<<<< SEARCH\n(.*?)\n=======\n(.*?)\n>>>>>>> REPLACE", re.DOTALL
)
self.prompt_pattern = re.compile(r"^> (.*)$", re.MULTILINE)
self.separator_pattern = re.compile(r"^─+$", re.MULTILINE)
self.token_info_pattern = re.compile(
r"Tokens: (.*?) sent, (.*?) received\. Cost: (.*?) message, (.*?) session\."
)
self.aider_pattern = re.compile(r"\baider\b", re.IGNORECASE)
def parse_output(self, output: str) -> List[AiderOutputBlock]:
"""Parse the output from the Aider model into blocks."""
blocks = []
current_block = []
for line in output.split("\n"):
line = line.rstrip("\r")
line = re.sub(self.aider_pattern, "Dev Assist", line)
if not line.strip():
if current_block:
blocks.append(self._create_text_block("".join(current_block)))
current_block = []
continue
if "<<<<<<< SEARCH" in line:
if current_block:
blocks.append(self._create_text_block("".join(current_block)))
current_block = []
blocks.append(self._create_code_block(line))
continue
if line.startswith("> "):
if current_block:
blocks.append(self._create_text_block("".join(current_block)))
current_block = []
prompt_content = line[2:].strip()
blocks.append(self._create_prompt_block(prompt_content))
continue
if re.match(self.separator_pattern, line):
if current_block:
blocks.append(self._create_text_block("".join(current_block)))
current_block = []
blocks.append(self._create_separator_block(line))
continue
if re.match(self.token_info_pattern, line):
if current_block:
blocks.append(self._create_text_block("".join(current_block)))
current_block = []
blocks.append(self._create_token_info_block(line))
continue
current_block.append(line + "\n")
if current_block:
blocks.append(self._create_text_block("".join(current_block)))
return blocks
def _create_text_block(self, content: str) -> AiderOutputBlock:
return AiderOutputBlock(content=content, type="text")
def _create_code_block(self, content: str) -> AiderOutputBlock:
return AiderOutputBlock(content=content, type="code", is_code_block=True)
def _create_prompt_block(self, content: str) -> AiderOutputBlock:
return AiderOutputBlock(content=content, type="prompt", is_prompt=True)
def _create_separator_block(self, content: str) -> AiderOutputBlock:
return AiderOutputBlock(content=content, type="separator", is_separator=True)
def _create_token_info_block(self, content: str) -> AiderOutputBlock:
return AiderOutputBlock(content=content, type="token_info", is_token_info=True)
|