from __future__ import annotations import hashlib import os import re import sys from dataclasses import dataclass from pathlib import Path from rich import print PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from scripts.common import ( CHUNKS_FILE, PAGES_ROOT, ZPWIKI_ROOT, load_zpwiki_page, write_json, ) MAX_TOKENS = int(os.getenv("CHUNK_MAX_TOKENS", "450")) OVERLAP_TOKENS = int(os.getenv("CHUNK_OVERLAP_TOKENS", "70")) MIN_CHUNK_TOKENS = int(os.getenv("CHUNK_MIN_TOKENS", "80")) TOKEN_ENCODING = os.getenv("CHUNK_TOKEN_ENCODING", "cl100k_base") HEADING_RE = re.compile(r"^(#{1,6})[ \t]+(.+?)[ \t]*#*[ \t]*$") FENCE_RE = re.compile(r"^[ \t]*(```+|~~~+)") LIST_ITEM_RE = re.compile(r"^[ \t]*(?:[-+*]|\d+[.)])[ \t]+") SENTENCE_SPLIT_RE = re.compile( r"(?<=[.!?…])(?:[\"'”’»\)\]]*)\s+" r"(?=[A-ZÁÄČĎÉÍĹĽŇÓÔŔŠŤÚÝŽ0-9])" ) FALLBACK_TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE) @dataclass(frozen=True) class Section: heading_path: tuple[str, ...] body: str @dataclass class ChunkDraft: context_lines: list[str] units: list[str] heading_paths: list[list[str]] @property def text(self) -> str: context = "\n".join(self.context_lines).strip() body = join_units(self.units) if context and body: return f"{context}\n\n{body}" return context or body class TokenCounter: """Použije tiktoken, prípadne jednoduchý lokálny odhad tokenov.""" def __init__(self, encoding_name: str) -> None: self.encoding_name = encoding_name self.encoding = None try: import tiktoken self.encoding = tiktoken.get_encoding(encoding_name) except (ImportError, ValueError): self.encoding = None @property def name(self) -> str: if self.encoding is None: return "regex fallback" return f"tiktoken/{self.encoding_name}" def count(self, text: str) -> int: if not text: return 0 if self.encoding is not None: return len( self.encoding.encode( text, disallowed_special=(), ) ) return len(FALLBACK_TOKEN_RE.findall(text)) TOKEN_COUNTER = TokenCounter(TOKEN_ENCODING) def clean_markdown(text: str) -> str: text = text.replace("\r\n", "\n").replace("\r", "\n") text = re.sub(r"[ \t]+\n", "\n", text) text = re.sub(r"\n{3,}", "\n\n", text) return text.strip() def normalize_label(text: str) -> str: text = text.casefold().strip() text = re.sub(r"\s+", " ", text) return text def parse_sections(text: str) -> list[Section]: """Rozdelí Markdown podľa nadpisov a zachová ich hierarchiu.""" lines = clean_markdown(text).splitlines() sections: list[Section] = [] heading_stack: list[str] = [] current_path: tuple[str, ...] = () body_lines: list[str] = [] active_fence: str | None = None def flush() -> None: body = "\n".join(body_lines).strip() # Prázdny rodičovský nadpis netvorí samostatný chunk. # Jeho názov zostane súčasťou heading_path podsekcií. if body: sections.append( Section( heading_path=current_path, body=body, ) ) body_lines.clear() for line in lines: fence_match = FENCE_RE.match(line) if fence_match: marker = fence_match.group(1)[0] if active_fence == marker: active_fence = None else: active_fence = marker body_lines.append(line) continue heading_match = None if active_fence else HEADING_RE.match(line) if not heading_match: body_lines.append(line) continue flush() level = len(heading_match.group(1)) heading = heading_match.group(2).strip() heading_stack[:] = heading_stack[: level - 1] heading_stack.append(heading) current_path = tuple(heading_stack) flush() return sections def split_markdown_blocks(text: str) -> list[str]: """Rozdelí text na odseky bez rozbitia fenced code blocku.""" blocks: list[str] = [] current: list[str] = [] active_fence: str | None = None def flush() -> None: block = "\n".join(current).strip() if block: blocks.append(block) current.clear() for line in text.splitlines(): fence_match = FENCE_RE.match(line) if active_fence is not None: current.append(line) if line.lstrip().startswith(active_fence): active_fence = None flush() continue if fence_match: flush() active_fence = fence_match.group(1) current.append(line) continue if not line.strip(): flush() else: current.append(line) flush() return blocks def is_code_block(text: str) -> bool: stripped = text.lstrip() return stripped.startswith("```") or stripped.startswith("~~~") def is_table(text: str) -> bool: lines = [ line.strip() for line in text.splitlines() if line.strip() ] if len(lines) < 2 or "|" not in lines[0]: return False return bool( re.fullmatch( r"\|?[ :|-]+\|?", lines[1], ) ) def split_sentences(text: str) -> list[str]: lines = [ line.strip() for line in text.splitlines() if line.strip() ] if lines and all(LIST_ITEM_RE.match(line) for line in lines): return lines normalized = " ".join(lines) if not normalized: return [] return [ part.strip() for part in SENTENCE_SPLIT_RE.split(normalized) if part.strip() ] def split_by_words(text: str, max_tokens: int) -> list[str]: """Posledná poistka: text rozdelí iba medzi slovami.""" words = re.findall(r"\S+", text) pieces: list[str] = [] current: list[str] = [] for word in words: candidate = " ".join([*current, word]) if current and TOKEN_COUNTER.count(candidate) > max_tokens: pieces.append(" ".join(current)) current = [word] else: current.append(word) if current: pieces.append(" ".join(current)) return pieces def split_code_block(text: str, max_tokens: int) -> list[str]: lines = text.splitlines() if len(lines) < 3: return split_by_words(text, max_tokens) opening = lines[0] has_closing = bool(FENCE_RE.match(lines[-1])) closing = lines[-1] if has_closing else "```" content = lines[1:-1] if has_closing else lines[1:] pieces: list[str] = [] current: list[str] = [] for line in content: candidate = "\n".join( [ opening, *current, line, closing, ] ) if current and TOKEN_COUNTER.count(candidate) > max_tokens: pieces.append( "\n".join( [ opening, *current, closing, ] ) ) current = [line] else: current.append(line) if current: pieces.append( "\n".join( [ opening, *current, closing, ] ) ) return pieces or [text] def split_table(text: str, max_tokens: int) -> list[str]: lines = [ line for line in text.splitlines() if line.strip() ] if len(lines) < 3: return [text] header = lines[:2] rows = lines[2:] pieces: list[str] = [] current_rows: list[str] = [] for row in rows: candidate = "\n".join( [ *header, *current_rows, row, ] ) if current_rows and TOKEN_COUNTER.count(candidate) > max_tokens: pieces.append( "\n".join( [ *header, *current_rows, ] ) ) current_rows = [row] else: current_rows.append(row) if current_rows: pieces.append( "\n".join( [ *header, *current_rows, ] ) ) return pieces or [text] def block_to_units(block: str, max_tokens: int) -> list[str]: """Vytvorí jednotky delené iba na prirodzených hraniciach.""" if is_code_block(block): if TOKEN_COUNTER.count(block) <= max_tokens: return [block] return split_code_block(block, max_tokens) if is_table(block): if TOKEN_COUNTER.count(block) <= max_tokens: return [block] return split_table(block, max_tokens) units: list[str] = [] for sentence in split_sentences(block): if TOKEN_COUNTER.count(sentence) <= max_tokens: units.append(sentence) else: units.extend( split_by_words( sentence, max_tokens, ) ) return units def join_units(units: list[str]) -> str: return "\n\n".join( unit.strip() for unit in units if unit.strip() ).strip() def unique_strings(values: list[str]) -> list[str]: result: list[str] = [] seen: set[str] = set() for value in values: key = normalize_label(value) if key and key not in seen: result.append(value) seen.add(key) return result def unique_paths(paths: list[list[str]]) -> list[list[str]]: result: list[list[str]] = [] seen: set[tuple[str, ...]] = set() for path in paths: if not path: continue key = tuple( normalize_label(item) for item in path ) if key not in seen: result.append(path) seen.add(key) return result def get_overlap( units: list[str], target_tokens: int, ) -> list[str]: """Vráti koncové celé vety alebo bloky pre ďalší chunk.""" if target_tokens <= 0 or len(units) <= 1: return [] selected: list[str] = [] for unit in reversed(units): candidate = [unit, *selected] if ( selected and TOKEN_COUNTER.count(join_units(candidate)) > target_tokens ): break selected = candidate if TOKEN_COUNTER.count(join_units(selected)) >= target_tokens: break # Celý chunk sa nesmie zopakovať ako overlap. if len(selected) == len(units): selected = selected[1:] return selected def pack_units( units: list[str], max_tokens: int, overlap_tokens: int, ) -> list[list[str]]: packed: list[list[str]] = [] current: list[str] = [] for unit in units: candidate = join_units([*current, unit]) if current and TOKEN_COUNTER.count(candidate) > max_tokens: packed.append(current.copy()) current = get_overlap( current, overlap_tokens, ) while ( current and TOKEN_COUNTER.count( join_units([*current, unit]) ) > max_tokens ): current.pop(0) current.append(unit) if current: packed.append(current.copy()) return packed def build_context_lines( document_title: str | None, heading_path: tuple[str, ...], ) -> list[str]: """Pridá názov dokumentu a podľa potreby aj cestu sekcie.""" title = (document_title or "").strip() path = [ item.strip() for item in heading_path if item.strip() ] lines: list[str] = [] if title: lines.append(f"Dokument: {title}") if path: visible_path = path # Ak je prvý Markdown nadpis rovnaký ako názov dokumentu, # v riadku Sekcia ho neopakujeme. if ( title and normalize_label(path[0]) == normalize_label(title) ): visible_path = path[1:] if visible_path: lines.append( f"Sekcia: {' > '.join(visible_path)}" ) return lines def heading_path_list( heading_path: tuple[str, ...], ) -> list[list[str]]: if not heading_path: return [] return [list(heading_path)] def chunk_section( section: Section, document_title: str | None, ) -> list[ChunkDraft]: context_lines = build_context_lines( document_title, section.heading_path, ) context_tokens = TOKEN_COUNTER.count( "\n".join(context_lines) ) body_limit = max( 1, MAX_TOKENS - context_tokens - 4, ) units: list[str] = [] for block in split_markdown_blocks(section.body): units.extend( block_to_units( block, body_limit, ) ) if not units: return [] packed_groups = pack_units( units, body_limit, OVERLAP_TOKENS, ) return [ ChunkDraft( context_lines=context_lines.copy(), units=packed_units, heading_paths=heading_path_list( section.heading_path ), ) for packed_units in packed_groups ] def merge_drafts( first: ChunkDraft, second: ChunkDraft, ) -> ChunkDraft: return ChunkDraft( context_lines=unique_strings( [ *first.context_lines, *second.context_lines, ] ), units=[ *first.units, *second.units, ], heading_paths=unique_paths( [ *first.heading_paths, *second.heading_paths, ] ), ) def token_count(draft: ChunkDraft) -> int: return TOKEN_COUNTER.count(draft.text) def try_merge_short_chunks( drafts: list[ChunkDraft], ) -> tuple[list[ChunkDraft], bool]: """Najprv skúsi krátky chunk úplne zlúčiť so susedom.""" result: list[ChunkDraft] = [] changed = False index = 0 while index < len(drafts): current = drafts[index] if ( token_count(current) < MIN_CHUNK_TOKENS and index + 1 < len(drafts) ): combined = merge_drafts( current, drafts[index + 1], ) if token_count(combined) <= MAX_TOKENS: result.append(combined) index += 2 changed = True continue if ( result and token_count(current) < MIN_CHUNK_TOKENS ): combined = merge_drafts( result[-1], current, ) if token_count(combined) <= MAX_TOKENS: result[-1] = combined index += 1 changed = True continue result.append(current) index += 1 return result, changed def try_balance_with_next( current: ChunkDraft, following: ChunkDraft, ) -> bool: """Presunie celé úvodné jednotky z ďalšieho chunku do krátkeho.""" changed = False while ( token_count(current) < MIN_CHUNK_TOKENS and len(following.units) > 1 ): moved_unit = following.units[0] candidate = ChunkDraft( context_lines=unique_strings( [ *current.context_lines, *following.context_lines, ] ), units=[ *current.units, moved_unit, ], heading_paths=unique_paths( [ *current.heading_paths, *following.heading_paths, ] ), ) remaining = ChunkDraft( context_lines=following.context_lines.copy(), units=following.units[1:], heading_paths=[ path.copy() for path in following.heading_paths ], ) if token_count(candidate) > MAX_TOKENS: break # Nevytvoríme nový krátky sused, # pokiaľ tým reálne nič nezlepšíme. if token_count(remaining) < MIN_CHUNK_TOKENS: break current.context_lines = candidate.context_lines current.units = candidate.units current.heading_paths = candidate.heading_paths following.units = remaining.units changed = True return changed def try_balance_with_previous( previous: ChunkDraft, current: ChunkDraft, ) -> bool: """Presunie celé koncové jednotky z predošlého chunku do krátkeho.""" changed = False while ( token_count(current) < MIN_CHUNK_TOKENS and len(previous.units) > 1 ): moved_unit = previous.units[-1] candidate = ChunkDraft( context_lines=unique_strings( [ *previous.context_lines, *current.context_lines, ] ), units=[ moved_unit, *current.units, ], heading_paths=unique_paths( [ *previous.heading_paths, *current.heading_paths, ] ), ) remaining = ChunkDraft( context_lines=previous.context_lines.copy(), units=previous.units[:-1], heading_paths=[ path.copy() for path in previous.heading_paths ], ) if token_count(candidate) > MAX_TOKENS: break if token_count(remaining) < MIN_CHUNK_TOKENS: break previous.units = remaining.units current.context_lines = candidate.context_lines current.units = candidate.units current.heading_paths = candidate.heading_paths changed = True return changed def balance_short_chunks( drafts: list[ChunkDraft], ) -> list[ChunkDraft]: """Zlúči alebo vyváži krátke chunky bez rezania viet a blokov.""" if len(drafts) < 2 or MIN_CHUNK_TOKENS <= 0: return drafts # Viac prechodov pomôže po zlúčení # susedných krátkych sekcií. for _ in range(3): drafts, merged = try_merge_short_chunks(drafts) balanced = False for index, current in enumerate(drafts): if token_count(current) >= MIN_CHUNK_TOKENS: continue if index + 1 < len(drafts): balanced = ( try_balance_with_next( current, drafts[index + 1], ) or balanced ) if ( token_count(current) < MIN_CHUNK_TOKENS and index > 0 ): balanced = ( try_balance_with_previous( drafts[index - 1], current, ) or balanced ) if not merged and not balanced: break return drafts def validate_settings() -> None: if MAX_TOKENS <= 0: raise ValueError( "CHUNK_MAX_TOKENS musí byť väčšie ako 0" ) if ( OVERLAP_TOKENS < 0 or OVERLAP_TOKENS >= MAX_TOKENS ): raise ValueError( "CHUNK_OVERLAP_TOKENS musí byť od 0 " "po MAX_TOKENS - 1" ) if ( MIN_CHUNK_TOKENS < 0 or MIN_CHUNK_TOKENS >= MAX_TOKENS ): raise ValueError( "CHUNK_MIN_TOKENS musí byť od 0 " "po MAX_TOKENS - 1" ) def chunk_markdown( text: str, document_title: str | None = None, ) -> list[ChunkDraft]: validate_settings() text = clean_markdown(text) if not text: return [] drafts: list[ChunkDraft] = [] for section in parse_sections(text): drafts.extend( chunk_section( section, document_title, ) ) return balance_short_chunks(drafts) def content_hash( document_path: str, text: str, ) -> str: value = f"{document_path}\n{text}".encode("utf-8") return hashlib.sha256(value).hexdigest()[:16] def build_chunks() -> list[dict]: if not PAGES_ROOT.exists(): raise SystemExit( f"Neexistuje priečinok: {PAGES_ROOT}" ) all_chunks: list[dict] = [] document_count = 0 for file_path in sorted( PAGES_ROOT.glob("**/README.md") ): document = load_zpwiki_page(file_path) document_count += 1 drafts = chunk_markdown( document["content"], document_title=document.get("title"), ) for index, draft in enumerate(drafts): text = draft.text token_total = TOKEN_COUNTER.count(text) all_chunks.append( { "chunk_id": ( f"{document['path']}::chunk-{index}" ), "document_path": document["path"], "title": document["title"], "categories": document["categories"], "tags": document["tags"], "author": document["author"], "published": document["published"], "chunk_index": index, "heading_paths": draft.heading_paths, "text": text, "text_length": len(text), "token_count": token_total, "content_hash": content_hash( document["path"], text, ), } ) write_json( CHUNKS_FILE, all_chunks, ) short_count = sum( chunk["token_count"] < MIN_CHUNK_TOKENS for chunk in all_chunks ) oversized_count = sum( chunk["token_count"] > MAX_TOKENS for chunk in all_chunks ) print( f"[green]ZPWIKI_ROOT:[/green] " f"{ZPWIKI_ROOT}" ) print( f"[green]Tokenizer:[/green] " f"{TOKEN_COUNTER.name}" ) print( f"[green]Dokumentov:[/green] " f"{document_count}" ) print( f"[green]Chunkov:[/green] " f"{len(all_chunks)}" ) print( f"[green]Krátkych chunkov:[/green] " f"{short_count}" ) print( f"[green]Príliš veľkých chunkov:[/green] " f"{oversized_count}" ) print( f"[green]Výstup uložený do:[/green] " f"{CHUNKS_FILE}" ) if all_chunks: print( "\n[bold]Ukážka prvého chunku:[/bold]" ) print(all_chunks[0]) return all_chunks def main() -> None: build_chunks() if __name__ == "__main__": main()