lambdagent/agentexample/qaagent67lambda/scripts/rebuild_index.py

#!/usr/bin/env python3
"""
Rebuild RAG index with improved chunking:
- Remove file headers and page markers
- Strict chunk size enforcement
- Better paragraph-aware splitting
- Filter out empty/noise chunks
"""
import json
import re
import math
from pathlib import Path
from collections import Counter

from config import cfg
PROCESSED_DIR = cfg.processed_dir
INDEX_FILE = cfg.index_file
KEYWORD_FILE = cfg.keyword_file

CHUNK_SIZE = 512
CHUNK_OVERLAP = 64
MIN_CHUNK_LEN = 30  # Discard chunks shorter than this


def clean_text(text, filename):
    """Remove noise from extracted PDF text and merge broken lines"""
    # Remove file header line
    text = re.sub(r'^# 源文件:.*\n*', '', text)
    # Remove page markers
    text = re.sub(r'--- 第\d+页 ---\n?', '', text)
    # Remove leading/trailing whitespace per line
    lines = [line.strip() for line in text.split('\n')]

    # Merge broken lines from PDF extraction.
    # PDF text has line breaks from page layout, not real paragraph breaks.
    # Strategy: only treat as paragraph break if previous text ends with
    # sentence-ending punctuation AND there's a blank line.
    merged = []
    buf = ''
    blank_count = 0
    for line in lines:
        if not line:
            blank_count += 1
            continue
        if buf:
            last_char = buf[-1] if buf else ''
            ends_sentence = last_char in '。！？；：）》」』"'
            # Real paragraph break = ended with punctuation + had blank line(s)
            if ends_sentence and blank_count > 0:
                merged.append(buf)
                merged.append('')  # preserve paragraph separator
                buf = line
            else:
                # Merge: either no blank line, or previous didn't end sentence
                buf += line
        else:
            buf = line
        blank_count = 0
    if buf:
        merged.append(buf)

    text = '\n'.join(merged)
    # Collapse multiple blank lines
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


def _smart_split(text, max_len, overlap):
    """Split long text at natural break points (punctuation), not mid-word"""
    # Priority: sentence-end > clause-end > comma > any punct > force
    break_chars = '。！？；\n'
    soft_break_chars = '，、：）》」』"'

    pieces = []
    start = 0
    while start < len(text):
        if start + max_len >= len(text):
            pieces.append(text[start:])
            break

        end = start + max_len
        # Look backwards for best break point
        best = -1
        # First try sentence-ending punctuation
        for i in range(end, max(start + max_len // 2, start), -1):
            if i < len(text) and text[i] in break_chars:
                best = i + 1
                break
        # If not found, try soft breaks (comma etc)
        if best == -1:
            for i in range(end, max(start + max_len // 2, start), -1):
                if i < len(text) and text[i] in soft_break_chars:
                    best = i + 1
                    break
        # Last resort: hard cut
        if best == -1:
            best = end

        pieces.append(text[start:best])
        # Overlap: go back a bit
        start = max(best - overlap, start + 1)
        if start >= len(text):
            break

    return pieces


def split_chunks(text, source_file):
    """Split text into chunks with strict size enforcement.
    Tables marked with [TABLE_START]...[TABLE_END] are kept as whole chunks."""
    chunks = []
    chunk_id = 0

    # First: extract tables as separate blocks
    # Split text into table blocks and text blocks
    parts = re.split(r'(\[TABLE_START\].*?\[TABLE_END\])', text, flags=re.DOTALL)

    all_paragraphs = []
    for part in parts:
        part = part.strip()
        if not part:
            continue
        if part.startswith('[TABLE_START]'):
            # Keep table as a single "paragraph" (won't be split further)
            table_content = part.replace('[TABLE_START]', '').replace('[TABLE_END]', '').strip()
            if table_content:
                all_paragraphs.append(('table', table_content))
        else:
            # Regular text: split into paragraphs
            for para in re.split(r'\n\s*\n', part):
                para = para.strip()
                if para and len(para) >= 5:
                    all_paragraphs.append(('text', para))

    # Now process paragraphs
    paragraphs_iter = all_paragraphs

    current_chunk = ''

    for ptype, para in paragraphs_iter:
        # Tables: always emit as their own chunk(s)
        if ptype == 'table':
            # Flush current text chunk first
            if current_chunk and len(current_chunk) >= MIN_CHUNK_LEN:
                chunks.append({
                    'id': f'{source_file}#chunk{chunk_id}',
                    'source': source_file,
                    'chunk_id': chunk_id,
                    'text': current_chunk,
                    'length': len(current_chunk),
                })
                chunk_id += 1
                current_chunk = ''

            # If table fits in one chunk, keep whole; otherwise split by rows
            if len(para) <= CHUNK_SIZE * 2:  # Allow tables up to 2x chunk size
                chunks.append({
                    'id': f'{source_file}#chunk{chunk_id}',
                    'source': source_file,
                    'chunk_id': chunk_id,
                    'text': para,
                    'length': len(para),
                })
                chunk_id += 1
            else:
                # Large table: split keeping header + groups of rows
                lines = para.split('\n')
                header = '\n'.join(lines[:2]) if len(lines) > 2 else ''  # header + separator
                table_chunk = header
                for line in lines[2:]:
                    if len(table_chunk) + len(line) + 1 > CHUNK_SIZE * 2:
                        if table_chunk:
                            chunks.append({
                                'id': f'{source_file}#chunk{chunk_id}',
                                'source': source_file,
                                'chunk_id': chunk_id,
                                'text': table_chunk,
                                'length': len(table_chunk),
                            })
                            chunk_id += 1
                        table_chunk = header + '\n' + line  # restart with header
                    else:
                        table_chunk += '\n' + line
                if table_chunk and len(table_chunk) >= MIN_CHUNK_LEN:
                    chunks.append({
                        'id': f'{source_file}#chunk{chunk_id}',
                        'source': source_file,
                        'chunk_id': chunk_id,
                        'text': table_chunk,
                        'length': len(table_chunk),
                    })
                    chunk_id += 1
            continue

        # Regular text paragraph
        # If adding this paragraph fits in chunk_size
        if len(current_chunk) + len(para) + 1 <= CHUNK_SIZE:
            current_chunk = (current_chunk + '\n' + para).strip() if current_chunk else para
        else:
            # Save current chunk if it has content
            if current_chunk and len(current_chunk) >= MIN_CHUNK_LEN:
                chunks.append({
                    'id': f'{source_file}#chunk{chunk_id}',
                    'source': source_file,
                    'chunk_id': chunk_id,
                    'text': current_chunk,
                    'length': len(current_chunk),
                })
                chunk_id += 1

            # If paragraph itself is too long, split by sentences/characters
            if len(para) > CHUNK_SIZE:
                # Try sentence splitting first (at major punctuation)
                sentences = re.split(r'(?<=[。！？；\n])', para)
                # If no split happened (single long sentence), try comma-level
                if len(sentences) <= 1 and len(para) > CHUNK_SIZE:
                    sentences = _smart_split(para, CHUNK_SIZE, CHUNK_OVERLAP)
                    # Already split into right-sized pieces, emit directly
                    for piece in sentences:
                        if len(piece.strip()) >= MIN_CHUNK_LEN:
                            chunks.append({
                                'id': f'{source_file}#chunk{chunk_id}',
                                'source': source_file,
                                'chunk_id': chunk_id,
                                'text': piece.strip(),
                                'length': len(piece.strip()),
                            })
                            chunk_id += 1
                    current_chunk = ''
                    continue
                sub_chunk = ''
                for sent in sentences:
                    sent = sent.strip()
                    if not sent:
                        continue
                    if len(sub_chunk) + len(sent) + 1 <= CHUNK_SIZE:
                        sub_chunk = (sub_chunk + sent).strip() if sub_chunk else sent
                    else:
                        if sub_chunk and len(sub_chunk) >= MIN_CHUNK_LEN:
                            chunks.append({
                                'id': f'{source_file}#chunk{chunk_id}',
                                'source': source_file,
                                'chunk_id': chunk_id,
                                'text': sub_chunk,
                                'length': len(sub_chunk),
                            })
                            chunk_id += 1
                        # If single sentence > CHUNK_SIZE, smart split at punctuation
                        if len(sent) > CHUNK_SIZE:
                            pieces = _smart_split(sent, CHUNK_SIZE, CHUNK_OVERLAP)
                            for piece in pieces:
                                if len(piece) >= MIN_CHUNK_LEN:
                                    chunks.append({
                                        'id': f'{source_file}#chunk{chunk_id}',
                                        'source': source_file,
                                        'chunk_id': chunk_id,
                                        'text': piece,
                                        'length': len(piece),
                                    })
                                    chunk_id += 1
                            sub_chunk = ''
                        else:
                            sub_chunk = sent
                # Don't forget remaining sub_chunk
                if sub_chunk and len(sub_chunk) >= MIN_CHUNK_LEN:
                    current_chunk = sub_chunk
                else:
                    current_chunk = ''
            else:
                # Start new chunk with overlap from previous
                overlap = current_chunk[-CHUNK_OVERLAP:] if len(current_chunk) > CHUNK_OVERLAP else ''
                current_chunk = (overlap + '\n' + para).strip() if overlap else para

    # Last chunk
    if current_chunk and len(current_chunk) >= MIN_CHUNK_LEN:
        chunks.append({
            'id': f'{source_file}#chunk{chunk_id}',
            'source': source_file,
            'chunk_id': chunk_id,
            'text': current_chunk,
            'length': len(current_chunk),
        })

    return chunks


# Build index
all_chunks = []
txt_files = sorted(PROCESSED_DIR.glob('*.txt'))
print(f'Processing {len(txt_files)} text files...')

for i, txt_path in enumerate(txt_files):
    text = txt_path.read_text(encoding='utf-8')
    text = clean_text(text, txt_path.stem)

    if len(text) < 30:
        continue

    chunks = split_chunks(text, txt_path.stem)
    all_chunks.extend(chunks)

    if (i + 1) % 200 == 0:
        print(f'  Progress: {i+1}/{len(txt_files)}, chunks: {len(all_chunks)}')

# Build keyword inverted index
print(f'Building keyword index... ({len(all_chunks)} chunks)')

keyword_index = {}
for chunk in all_chunks:
    words = set()
    text = chunk['text'].lower()
    # Chinese bigram + trigram
    chinese_chars = re.findall(r'[\u4e00-\u9fff]+', text)
    for cc in chinese_chars:
        for j in range(len(cc) - 1):
            words.add(cc[j:j+2])
        if len(cc) >= 3:
            for j in range(len(cc) - 2):
                words.add(cc[j:j+3])
    # English words
    eng_words = re.findall(r'[a-z]+', text)
    words.update(eng_words)

    for w in words:
        if w not in keyword_index:
            keyword_index[w] = []
        keyword_index[w].append(chunk['id'])

# Save
index_data = {
    'total_chunks': len(all_chunks),
    'total_files': len(txt_files),
    'chunk_size': CHUNK_SIZE,
    'chunks': all_chunks,
    'keyword_index_size': len(keyword_index),
}

with open(INDEX_FILE, 'w', encoding='utf-8') as f:
    json.dump(index_data, f, ensure_ascii=False)

with open(KEYWORD_FILE, 'w', encoding='utf-8') as f:
    json.dump(keyword_index, f, ensure_ascii=False)

# Stats
lengths = [c['length'] for c in all_chunks]
print(f'\nDone!')
print(f'  Total chunks: {len(all_chunks)}')
print(f'  Keywords: {len(keyword_index)}')
print(f'  Avg length: {sum(lengths)/len(lengths):.0f}')
print(f'  Min/Max: {min(lengths)}/{max(lengths)}')
print(f'  <100: {sum(1 for l in lengths if l < 100)} ({sum(1 for l in lengths if l < 100)*100/len(lengths):.1f}%)')
print(f'  >1000: {sum(1 for l in lengths if l > 1000)} ({sum(1 for l in lengths if l > 1000)*100/len(lengths):.1f}%)')
print(f'  Index size: {INDEX_FILE.stat().st_size/1024/1024:.1f} MB')