Harshavard21/FinRAG

1"""2src/generation/response_parser.py3===================================4Parses LLM responses to extract structured citation info.5"""6 7from __future__ import annotations8import re9 10 11def extract_citations(answer_text: str) -> list[dict]:12    """13    Extract [SOURCE N] citation markers from LLM answer text.14 15    Returns list of dicts with citation info parsed from the answer.16    """17    citations = []18    pattern = r'\[SOURCE\s+(\d+)\]'19    matches = re.finditer(pattern, answer_text)20    for match in matches:21        idx = int(match.group(1))22        if not any(c["index"] == idx for c in citations):23            citations.append({"index": idx, "marker": match.group(0)})24    return citations25 26 27def clean_answer(answer_text: str) -> str:28    """Remove redundant whitespace from LLM output."""29    lines = answer_text.split("\n")30    cleaned = []31    prev_blank = False32    for line in lines:33        is_blank = not line.strip()34        if is_blank and prev_blank:35            continue36        cleaned.append(line)37        prev_blank = is_blank38    return "\n".join(cleaned).strip()39