diff --git a/rag/chunk.py b/rag/chunk.py index b8d7317..81ef39c 100644 --- a/rag/chunk.py +++ b/rag/chunk.py @@ -31,6 +31,27 @@ from typing import Iterator CHARS_PER_TOKEN = 4 TARGET_TOKENS = 500 TARGET_CHARS = TARGET_TOKENS * CHARS_PER_TOKEN +# Hard cap: nomic-embed-text's context is 2048 tokens. Anything larger +# 400s the entire embed batch. 6000 chars ≈ 1500 tokens leaves headroom. +MAX_CHARS = 6000 + + +def _hard_split(text: str) -> list[str]: + """Split an oversized block on line boundaries into MAX_CHARS pieces.""" + if len(text) <= MAX_CHARS: + return [text] + out: list[str] = [] + buf: list[str] = [] + buf_chars = 0 + for line in text.splitlines(keepends=True): + if buf_chars + len(line) > MAX_CHARS and buf: + out.append("".join(buf).rstrip()) + buf, buf_chars = [], 0 + buf.append(line) + buf_chars += len(line) + if buf: + out.append("".join(buf).rstrip()) + return out def estimate_tokens(text: str) -> int: @@ -104,23 +125,26 @@ def chunks_from_page( # ----- Body chunks: pack paragraphs up to TARGET_CHARS ------- ordinal = 1 + + def emit(buf: list[str]) -> Iterator[dict]: + nonlocal ordinal + merged = "\n\n".join(buf) + for piece in _hard_split(merged): + yield { + "id": f"{metadata['bundle_id']}::{page_id}::{ordinal}", + "text": piece, + "metadata": {**metadata, "ordinal": ordinal}, + } + ordinal += 1 + buf: list[str] = [] buf_chars = 0 for p in paragraphs: if buf_chars + len(p) > TARGET_CHARS and buf: - yield { - "id": f"{metadata['bundle_id']}::{page_id}::{ordinal}", - "text": "\n\n".join(buf), - "metadata": {**metadata, "ordinal": ordinal}, - } - ordinal += 1 + yield from emit(buf) buf = [] buf_chars = 0 buf.append(p) buf_chars += len(p) if buf: - yield { - "id": f"{metadata['bundle_id']}::{page_id}::{ordinal}", - "text": "\n\n".join(buf), - "metadata": {**metadata, "ordinal": ordinal}, - } + yield from emit(buf) diff --git a/rag/index.py b/rag/index.py index 8d1c74f..c512e80 100644 --- a/rag/index.py +++ b/rag/index.py @@ -29,7 +29,7 @@ CHROMA_DIR = ROOT / "chroma" # Collection name — convention: _docs. Override via env if needed. import os -PRODUCT_NAME = os.environ.get("PRODUCT_NAME", "myproduct") +PRODUCT_NAME = os.environ.get("PRODUCT_NAME", "hvm") COLLECTION = f"{PRODUCT_NAME}_docs"