mindnet/app/core/chunking/chunking_strategies.py

"""
FILE: app/core/chunking/chunking_strategies.py
DESCRIPTION: Strategien für atomares Sektions-Chunking (WP-15b konform).
"""
from typing import List, Dict, Any, Optional
from .chunking_models import RawBlock, Chunk
from .chunking_utils import estimate_tokens
from .chunking_parser import split_sentences

def _create_context_win(doc_title: str, sec_title: Optional[str], text: str) -> str:
    """Baut den Breadcrumb-Kontext für das Embedding-Fenster."""
    parts = []
    if doc_title: parts.append(doc_title)
    if sec_title and sec_title != doc_title: parts.append(sec_title)
    prefix = " > ".join(parts)
    return f"{prefix}\n{text}".strip() if prefix else text

def strategy_by_heading(blocks: List[RawBlock], config: Dict[str, Any], note_id: str, doc_title: str = "") -> List[Chunk]:
    """
    Gruppiert Blöcke zu Sektionen und hält diese atomar zusammen.
    """
    strict = config.get("strict_heading_split", False)
    target = config.get("target", 400)
    max_tokens = config.get("max", 600)
    split_level = config.get("split_level", 2)
    overlap = sum(config.get("overlap", (50, 80))) // 2
    
    chunks: List[Chunk] = []
    buf: List[RawBlock] = []
    cur_tokens = 0

    def _add_to_chunks(txt, title, path):
        idx = len(chunks)
        win = _create_context_win(doc_title, title, txt)
        chunks.append(Chunk(
            id=f"{note_id}#c{idx:02d}", note_id=note_id, index=idx,
            text=txt, window=win, token_count=estimate_tokens(txt),
            section_title=title, section_path=path,
            neighbors_prev=None, neighbors_next=None
        ))

    def _flush_buffer():
        nonlocal buf, cur_tokens
        if not buf: return
        
        main_title = buf[0].section_title
        main_path = buf[0].section_path
        full_text = "\n\n".join([b.text for b in buf])
        
        if estimate_tokens(full_text) <= max_tokens:
            _add_to_chunks(full_text, main_title, main_path)
        else:
            # Nur wenn eine Sektion ALLEINE zu groß ist, wird intern gesplittet
            sents = split_sentences(full_text)
            cur_sents = []; sub_len = 0
            for s in sents:
                slen = estimate_tokens(s)
                if sub_len + slen > target and cur_sents:
                    _add_to_chunks(" ".join(cur_sents), main_title, main_path)
                    # Overlap Logic
                    ov_s = []; ov_l = 0
                    for os in reversed(cur_sents):
                        if ov_l + estimate_tokens(os) < overlap:
                            ov_s.insert(0, os); ov_l += estimate_tokens(os)
                        else: break
                    cur_sents = list(ov_s); cur_sents.append(s); sub_len = ov_l + slen
                else: cur_sents.append(s); sub_len += slen
            if cur_sents: _add_to_chunks(" ".join(cur_sents), main_title, main_path)
        
        buf = []; cur_tokens = 0

    # SCHRITT 1: Gruppierung in atomare Sektions-Einheiten (Heading + Paragraphs)
    sections: List[List[RawBlock]] = []
    curr_sec: List[RawBlock] = []
    
    for b in blocks:
        # Ein Split-Trigger startet eine neue Sektion
        if b.kind == "heading" and b.level <= split_level:
            if curr_sec: sections.append(curr_sec)
            curr_sec = [b]
        else:
            curr_sec.append(b)
    if curr_sec: sections.append(curr_sec)

    # SCHRITT 2: Verarbeitung der Sektionen mit Vorausschau
    for sec in sections:
        # Token-Schätzung für die gesamte Sektion inkl. Newline-Overhead
        sec_text = "\n\n".join([b.text for b in sec])
        sec_tokens = estimate_tokens(sec_text)
        
        if buf:
            # Passt die Sektion noch in den aktuellen Chunk?
            if cur_tokens + sec_tokens > max_tokens:
                _flush_buffer()
            # Wenn strict: Jede neue Sektion auf split_level erzwingt neuen Chunk
            elif strict and sec[0].kind == "heading" and sec[0].level == split_level:
                _flush_buffer()
            # Wenn target erreicht: Neue Sektion startet neuen Chunk
            elif cur_tokens >= target:
                _flush_buffer()
        
        buf.extend(sec)
        cur_tokens += sec_tokens
        
        # Falls der Puffer (selbst nach flush) durch eine Riesen-Sektion zu groß ist
        if cur_tokens >= max_tokens:
            _flush_buffer()

    _flush_buffer()
    return chunks

def strategy_sliding_window(blocks: List[RawBlock], config: Dict[str, Any], note_id: str, context_prefix: str = "") -> List[Chunk]:
    # (Identische Korrektur wie oben für Sliding Window, falls benötigt)
    # Hier halten wir es einfach: Blöcke nacheinander bis target.
    target = config.get("target", 400)
    max_tokens = config.get("max", 600)
    
    chunks: List[Chunk] = []
    buf: List[RawBlock] = []
    
    for b in blocks:
        b_tokens = estimate_tokens(b.text)
        current_tokens = estimate_tokens("\n\n".join([x.text for x in buf])) if buf else 0
        
        if current_tokens + b_tokens > max_tokens and buf:
            txt = "\n\n".join([x.text for x in buf])
            idx = len(chunks)
            win = f"{context_prefix}\n{txt}".strip() if context_prefix else txt
            chunks.append(Chunk(id=f"{note_id}#c{idx:02d}", note_id=note_id, index=idx, text=txt, window=win, token_count=current_tokens, section_title=buf[0].section_title, section_path=buf[0].section_path, neighbors_prev=None, neighbors_next=None))
            buf = []
            current_tokens = 0
            
        buf.append(b)
        
    if buf:
        txt = "\n\n".join([x.text for x in buf])
        idx = len(chunks)
        win = f"{context_prefix}\n{txt}".strip() if context_prefix else txt
        chunks.append(Chunk(id=f"{note_id}#c{idx:02d}", note_id=note_id, index=idx, text=txt, window=win, token_count=estimate_tokens(txt), section_title=buf[0].section_title, section_path=buf[0].section_path, neighbors_prev=None, neighbors_next=None))
        
    return chunks