Spaces:
Sleeping
Sleeping
| """Module to parse the output of the Aider model.""" | |
| import re | |
| from dataclasses import dataclass | |
| from typing import List | |
| 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) | |