| 7 | |
| 8 | |
| 9 | class AnswerExtractor: |
| 10 | def __init__(self, config: AnswerExtractionConfig) -> None: |
| 11 | self.config = config |
| 12 | |
| 13 | def format_prompts(self, base_instructions: Union[str, List[str]]) -> Union[str, List[str]]: |
| 14 | """ |
| 15 | Combine base instruction with answer extraction instruction. |
| 16 | |
| 17 | Args: |
| 18 | base_instruction: Task-specific output instruction |
| 19 | |
| 20 | Returns: |
| 21 | Complete instruction with answer format |
| 22 | """ |
| 23 | if self.config is None or not self.config.enabled or not self.config.instruction: |
| 24 | return base_instructions |
| 25 | else: |
| 26 | is_single_instruction = isinstance(base_instructions, str) |
| 27 | base_instructions = [base_instructions] if is_single_instruction else base_instructions |
| 28 | instructions: List[str] = [] |
| 29 | for base_instruction in base_instructions: |
| 30 | separator = '' if not base_instruction else ' ' |
| 31 | instructions.append(f"{base_instruction}{separator}{self.config.instruction}") |
| 32 | return instructions[0] if is_single_instruction else instructions |
| 33 | |
| 34 | def extract_answers(self, response: Union[str, List[str], List[List[str]]]) -> Optional[Union[str, List[str], List[List[str]]]]: |
| 35 | """ |
| 36 | Extract answer from LLM response. |
| 37 | |
| 38 | Supports two extraction modes: |
| 39 | 1. XML-style tags (e.g., "<answer>"): Extracts content between paired tags |
| 40 | - If closing tag found: extracts between <answer> and </answer> |
| 41 | - If no closing tag: extracts everything after <answer> |
| 42 | 2. Simple markers (e.g., "####"): Extracts everything after the marker |
| 43 | |
| 44 | Args: |
| 45 | response: LLM output text |
| 46 | |
| 47 | Returns: |
| 48 | Extracted answer or None if tag not found |
| 49 | """ |
| 50 | if self.config is None or not self.config.enabled: |
| 51 | return response # Return full response if extraction disabled |
| 52 | |
| 53 | def extract_answer_per_response(tag: str, response: str) -> Optional[str]: |
| 54 | if tag not in response: |
| 55 | return None |
| 56 | |
| 57 | # Check if tag is XML-style (starts with < and ends with >) |
| 58 | if tag.startswith("<") and tag.endswith(">") and len(tag) > 2: |
| 59 | # Extract tag name (e.g., "answer" from "<answer>") |
| 60 | tag_name = tag[1:-1] |
| 61 | closing_tag = f"</{tag_name}>" |
| 62 | |
| 63 | # Try to find content between opening and closing tags |
| 64 | start_idx = response.find(tag) |
| 65 | if start_idx != -1: |
| 66 | content_start = start_idx + len(tag) |