dp-zp-agent/scripts/build_chunks.py
2026-07-28 23:56:03 +02:00

1029 lines
23 KiB
Python

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()