chromerag 0.1.0__py3-none-any.whl

chromerag/__init__.py

@@ -0,0 +1,37 @@
+"""ChromeRAG — Ingest-time elimination of site template noise for enterprise web RAG.
+
+Core contract: HTML string in → RAG-ready Markdown out.
+Fetching/crawling is intentionally out of scope (use your crawler / poc harness).
+"""
+
+from chromerag.batch import learn_then_extract
+from chromerag.config import ContentPriority, PageType, PipelineConfig, Strictness
+from chromerag.extractor import ChromeRAG
+from chromerag.input_quality import InputQualityReport, assess_input_html
+from chromerag.models import ExtractResult
+from chromerag.site_chrome import (
+    SiteChromeModel,
+    load_chrome_models,
+    mine_site_chrome,
+    save_chrome_models,
+)
+
+__all__ = [
+    "ChromeRAG",
+    "ExtractResult",
+    "PipelineConfig",
+    "ContentPriority",
+    "Strictness",
+    "PageType",
+    "SiteChromeModel",
+    "mine_site_chrome",
+    "save_chrome_models",
+    "load_chrome_models",
+    "learn_then_extract",
+    "assess_input_html",
+    "InputQualityReport",
+]
+__version__ = "0.1.0"
+__paper__ = (
+    "ChromeRAG: Ingest-Time Elimination of Site Template Noise for Enterprise Web RAG"
+)

chromerag/batch.py

@@ -0,0 +1,111 @@
+"""Learn-then-extract batch helpers.
+
+Critical order when STCE is enabled:
+  1. Mine chrome patterns across the batch (or load a saved model)
+  2. Only then extract each page with the learned model
+
+Never extract-first then learn later — early chunks would keep template noise.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable
+
+from chromerag.config import PipelineConfig
+from chromerag.extractor import ChromeRAG
+from chromerag.models import ExtractResult
+from chromerag.site_chrome import (
+    SiteChromeModel,
+    mine_site_chrome,
+    select_model_for_url,
+)
+
+
+@dataclass
+class BatchPage:
+    html: str
+    url: str | None = None
+    page_id: str | None = None
+
+
+@dataclass
+class BatchResult:
+    page_id: str | None
+    url: str | None
+    result: ExtractResult
+    chrome_group: str | None
+    stce_applied: bool
+
+
+def learn_chrome(
+    pages: Iterable[BatchPage | tuple[str | None, str]],
+    *,
+    min_pages: int | None = None,
+    frequency: float | None = None,
+    max_block_chars: int | None = None,
+    require_chrome_features: bool = True,
+    path_prefix_depth: int = 1,
+    config: PipelineConfig | None = None,
+) -> dict[str, SiteChromeModel]:
+    """Mine site chrome from a multi-page corpus. Does not extract."""
+    cfg = config or PipelineConfig()
+    pairs: list[tuple[str | None, str]] = []
+    for p in pages:
+        if isinstance(p, BatchPage):
+            pairs.append((p.url, p.html))
+        else:
+            pairs.append((p[0], p[1]))
+    return mine_site_chrome(
+        pairs,
+        min_pages=min_pages if min_pages is not None else cfg.stce_min_pages,
+        frequency=frequency if frequency is not None else cfg.stce_frequency,
+        max_block_chars=max_block_chars
+        if max_block_chars is not None
+        else cfg.stce_max_block_chars,
+        require_chrome_features=require_chrome_features,
+        path_prefix_depth=path_prefix_depth,
+    )
+
+
+def learn_then_extract(
+    pages: Iterable[BatchPage],
+    *,
+    config: PipelineConfig | None = None,
+    chrome_models: dict[str, SiteChromeModel] | None = None,
+    learn: bool = True,
+    path_prefix_depth: int = 1,
+) -> tuple[list[BatchResult], dict[str, SiteChromeModel]]:
+    """
+    If learn=True and chrome_models is None: mine from `pages` first, then extract.
+    If chrome_models provided: use them (skip mining).
+    If learn=False and no models: extract with heuristics/DVDF only (STCE off).
+    """
+    page_list = list(pages)
+    cfg = config or PipelineConfig()
+
+    models: dict[str, SiteChromeModel] = {}
+    if chrome_models is not None:
+        models = chrome_models
+    elif learn and cfg.enable_stce:
+        models = learn_chrome(page_list, config=cfg, path_prefix_depth=path_prefix_depth)
+
+    results: list[BatchResult] = []
+    for page in page_list:
+        model = (
+            select_model_for_url(models, page.url, path_prefix_depth=path_prefix_depth)
+            if models and cfg.enable_stce
+            else None
+        )
+        extractor = ChromeRAG(config=cfg, site_chrome=model)
+        result = extractor.extract(page.html, url=page.url)
+        results.append(
+            BatchResult(
+                page_id=page.page_id,
+                url=page.url,
+                result=result,
+                chrome_group=model.group_key if model else None,
+                stce_applied=bool(model and model.signatures),
+            )
+        )
+    return results, models

chromerag/cli.py

@@ -0,0 +1,296 @@
+"""ChromeRAG CLI — HTML → RAG Markdown (fetch is not a core concern).
+
+Subcommands:
+  extract  — single HTML file → markdown
+  learn    — mine site-chrome model from a directory of HTML fixtures
+  batch    — learn-then-extract (or extract with a saved chrome model)
+
+URL fetching is intentionally omitted from the default path. Use `poc.fetch_pages`
+or your own crawler, then point ChromeRAG at the saved HTML.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from chromerag.batch import BatchPage, learn_chrome, learn_then_extract
+from chromerag.config import ContentPriority, PipelineConfig, Strictness
+from chromerag.extractor import ChromeRAG
+from chromerag.site_chrome import load_chrome_models, save_chrome_models
+
+
+def _build_config(args: argparse.Namespace) -> PipelineConfig:
+    overrides = dict(
+        enable_dvdf=not getattr(args, "no_dvdf", False),
+        enable_schema=not getattr(args, "no_schema", False),
+        enable_tables=not getattr(args, "no_tables", False),
+        enable_stce=not getattr(args, "no_stce", False),
+        inject_heading_paths=getattr(args, "heading_paths", False),
+    )
+    if getattr(args, "priority", None):
+        return PipelineConfig.from_priority(args.priority, **overrides)
+    if getattr(args, "strictness", None):
+        return PipelineConfig.from_strictness(args.strictness, **overrides)
+    return PipelineConfig.from_priority(ContentPriority.BALANCED, **overrides)
+
+
+def _add_shared_flags(p: argparse.ArgumentParser) -> None:
+    p.add_argument(
+        "--priority",
+        choices=[c.value for c in ContentPriority],
+        default=None,
+        help="precision | coverage | balanced (default)",
+    )
+    p.add_argument(
+        "--strictness",
+        choices=[s.value for s in Strictness],
+        default=None,
+        help="Legacy alias for priority",
+    )
+    p.add_argument("--no-dvdf", action="store_true")
+    p.add_argument("--no-schema", action="store_true")
+    p.add_argument("--no-tables", action="store_true")
+    p.add_argument("--no-stce", action="store_true")
+    p.add_argument("--heading-paths", action="store_true")
+
+
+def _load_fixture_dir(raw_dir: Path) -> list[BatchPage]:
+    pages: list[BatchPage] = []
+    for meta_path in sorted(raw_dir.glob("*.meta.json")):
+        if meta_path.name.startswith("_"):
+            continue
+        meta = json.loads(meta_path.read_text(encoding="utf-8"))
+        html_path = raw_dir / f"{meta['id']}.html"
+        if not html_path.exists():
+            continue
+        pages.append(
+            BatchPage(
+                html=html_path.read_text(encoding="utf-8", errors="ignore"),
+                url=meta.get("url"),
+                page_id=meta["id"],
+            )
+        )
+    if not pages:
+        # Fallback: bare *.html without meta
+        for html_path in sorted(raw_dir.glob("*.html")):
+            pages.append(
+                BatchPage(
+                    html=html_path.read_text(encoding="utf-8", errors="ignore"),
+                    url=None,
+                    page_id=html_path.stem,
+                )
+            )
+    return pages
+
+
+def cmd_extract(args: argparse.Namespace) -> int:
+    path = Path(args.html)
+    if not path.exists():
+        sys.stderr.write(
+            f"File not found: {path}\n"
+            "ChromeRAG extracts HTML files only. Fetch pages with your crawler "
+            "or: python -m poc.fetch_pages\n"
+        )
+        return 2
+    html = path.read_text(encoding="utf-8", errors="ignore")
+    cfg = _build_config(args)
+    model = None
+    if args.chrome_model and not args.no_stce:
+        models = load_chrome_models(args.chrome_model)
+        from chromerag.site_chrome import select_model_for_url
+
+        model = select_model_for_url(models, args.url)
+    extractor = ChromeRAG(config=cfg, site_chrome=model)
+    result = extractor.extract(html, url=args.url)
+
+    for warning in result.warnings:
+        sys.stderr.write(f"WARNING: {warning}\n")
+
+    if args.output:
+        Path(args.output).write_text(result.markdown, encoding="utf-8")
+    else:
+        sys.stdout.write(result.markdown)
+
+    if args.json_meta:
+        meta = {
+            "url": result.url,
+            "tokens_estimate": result.tokens_estimate,
+            "n_blocks_in": result.n_blocks_in,
+            "n_blocks_kept": result.n_blocks_kept,
+            "n_tables": result.n_tables,
+            "front_matter": result.front_matter,
+            "warnings": result.warnings,
+            "input_quality": result.input_quality,
+            "diagnostics": result.diagnostics,
+        }
+        sys.stderr.write(json.dumps(meta, indent=2) + "\n")
+    # Non-zero only on hard failures; thin JS shells warn but still exit 0 so
+    # pipelines can log and decide. Use --fail-on-thin to hard-fail.
+    if getattr(args, "fail_on_thin", False) and result.input_quality.get("is_thin"):
+        return 3
+    return 0
+
+
+def cmd_learn(args: argparse.Namespace) -> int:
+    pages = _load_fixture_dir(Path(args.html_dir))
+    if len(pages) < 2:
+        sys.stderr.write("Need at least 2 HTML pages to learn chrome patterns.\n")
+        return 2
+    cfg = _build_config(args)
+    models = learn_chrome(
+        pages,
+        config=cfg,
+        min_pages=args.min_pages,
+        path_prefix_depth=args.path_prefix_depth,
+    )
+    out = Path(args.output)
+    save_chrome_models(models, out)
+    active = {k: len(m.signatures) for k, m in models.items() if m.signatures}
+    sys.stderr.write(
+        json.dumps(
+            {
+                "pages": len(pages),
+                "groups": {k: m.n_pages for k, m in models.items()},
+                "chrome_signatures": active,
+                "wrote": str(out),
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+    return 0
+
+
+def cmd_batch(args: argparse.Namespace) -> int:
+    """Learn-then-extract (default) or extract with a pre-learned chrome model."""
+    pages = _load_fixture_dir(Path(args.html_dir))
+    if not pages:
+        sys.stderr.write(f"No HTML fixtures in {args.html_dir}\n")
+        return 2
+    cfg = _build_config(args)
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    chrome_models = None
+    learn = True
+    if args.chrome_model:
+        chrome_models = load_chrome_models(args.chrome_model)
+        learn = False
+    if args.no_learn:
+        learn = False
+
+    results, models = learn_then_extract(
+        pages,
+        config=cfg,
+        chrome_models=chrome_models,
+        learn=learn,
+        path_prefix_depth=args.path_prefix_depth,
+    )
+
+    if args.save_chrome:
+        save_chrome_models(models, args.save_chrome)
+
+    summary = []
+    for br in results:
+        pid = br.page_id or "page"
+        page_out = out_dir / pid
+        page_out.mkdir(parents=True, exist_ok=True)
+        (page_out / "chromerag.md").write_text(br.result.markdown, encoding="utf-8")
+        # UI compat alias
+        (page_out / "smart_rag_extractor.md").write_text(br.result.markdown, encoding="utf-8")
+        iq = br.result.input_quality or {}
+        row = {
+            "id": pid,
+            "url": br.url,
+            "tokens": br.result.tokens_estimate,
+            "stce_applied": br.stce_applied,
+            "chrome_group": br.chrome_group,
+            "stce_removed": br.result.diagnostics.get("stce_removed", 0),
+            "priority": br.result.diagnostics.get("content_priority"),
+            "json_ld": br.result.diagnostics.get("schema_sources", {}).get("json_ld_count", 0),
+            "input_thin": bool(iq.get("is_thin")),
+            "likely_js_shell": bool(iq.get("likely_js_shell")),
+            "warnings": list(br.result.warnings or []),
+        }
+        summary.append(row)
+        flag = " THIN" if row["input_thin"] else ""
+        if row["likely_js_shell"]:
+            flag = " JS_SHELL"
+        print(
+            f"{pid:28s} tok={row['tokens']:6d} stce={row['stce_applied']} "
+            f"rm={row['stce_removed']} jsonld={row['json_ld']}{flag}"
+        )
+        for warning in br.result.warnings or []:
+            sys.stderr.write(f"WARNING [{pid}]: {warning}\n")
+
+    (out_dir / "batch_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="chromerag",
+        description=(
+            "ChromeRAG: ingest-time elimination of site template noise "
+            "for enterprise web RAG. Pass HTML files — do not use as a crawler."
+        ),
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_ex = sub.add_parser("extract", help="Extract one HTML file to Markdown")
+    p_ex.add_argument("html", help="Path to .html file")
+    p_ex.add_argument("-o", "--output", help="Write markdown to path")
+    p_ex.add_argument("--url", help="Canonical URL (for page-type + STCE group)")
+    p_ex.add_argument("--chrome-model", help="Path to learned chrome JSON from `chromerag learn`")
+    p_ex.add_argument("--json-meta", action="store_true")
+    p_ex.add_argument(
+        "--fail-on-thin",
+        action="store_true",
+        help="Exit 3 when input HTML looks thin / like a JS shell (still prints WARNINGs by default)",
+    )
+    _add_shared_flags(p_ex)
+    p_ex.set_defaults(func=cmd_extract)
+
+    p_learn = sub.add_parser(
+        "learn",
+        help="Mine site-chrome patterns from a directory of HTML (run BEFORE batch extract)",
+    )
+    p_learn.add_argument("html_dir", help="Directory with *.html (+ optional *.meta.json)")
+    p_learn.add_argument("-o", "--output", required=True, help="Write chrome model JSON here")
+    p_learn.add_argument("--min-pages", type=int, default=3)
+    p_learn.add_argument("--path-prefix-depth", type=int, default=1)
+    _add_shared_flags(p_learn)
+    p_learn.set_defaults(func=cmd_learn)
+
+    p_batch = sub.add_parser(
+        "batch",
+        help="Learn-then-extract a directory (or extract with --chrome-model / --no-learn)",
+    )
+    p_batch.add_argument("html_dir", help="Directory with HTML fixtures")
+    p_batch.add_argument("-o", "--output-dir", required=True, help="Output directory")
+    p_batch.add_argument(
+        "--chrome-model",
+        help="Use a previously learned chrome model (skips mining)",
+    )
+    p_batch.add_argument(
+        "--save-chrome",
+        help="Also write the mined chrome model to this path",
+    )
+    p_batch.add_argument(
+        "--no-learn",
+        action="store_true",
+        help="Skip STCE mining; heuristics + DVDF only",
+    )
+    p_batch.add_argument("--path-prefix-depth", type=int, default=1)
+    _add_shared_flags(p_batch)
+    p_batch.set_defaults(func=cmd_batch)
+
+    args = parser.parse_args(argv)
+    return int(args.func(args))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

chromerag/config.py

@@ -0,0 +1,181 @@
+"""Pipeline configuration — manageable strictness knobs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from chromerag.density import DensityConfig
+
+
+class ContentPriority(str, Enum):
+    """User-facing tradeoff: clean vectors vs keeping more relevant text.
+
+    PRECISION — Prefer correctness of extracted text. Aggressively drop chrome;
+                may lose some relevant chunks (lower content recall).
+    COVERAGE  — Prefer not losing relevant chunks. Tolerates more residual noise
+                (higher content recall, higher noise retention).
+    BALANCED  — Default enterprise RAG compromise.
+    """
+
+    PRECISION = "precision"
+    COVERAGE = "coverage"
+    BALANCED = "balanced"
+
+
+class Strictness(str, Enum):
+    """Internal aggressiveness presets (mapped from ContentPriority)."""
+
+    LENIENT = "lenient"  # Prefer recall — keep more text
+    BALANCED = "balanced"  # Default for enterprise RAG
+    AGGRESSIVE = "aggressive"  # Prefer precision — more chrome gone
+
+
+class PageType(str, Enum):
+    DOCS = "docs"
+    PRICING = "pricing"
+    MARKETING = "marketing"
+    HUB = "hub"
+    UNKNOWN = "unknown"
+
+
+_PRIORITY_TO_STRICTNESS = {
+    ContentPriority.PRECISION: Strictness.AGGRESSIVE,
+    ContentPriority.COVERAGE: Strictness.LENIENT,
+    ContentPriority.BALANCED: Strictness.BALANCED,
+}
+
+
+def priority_to_strictness(priority: ContentPriority | str) -> Strictness:
+    if isinstance(priority, str):
+        priority = ContentPriority(priority)
+    return _PRIORITY_TO_STRICTNESS[priority]
+
+
+@dataclass
+class PipelineConfig:
+    """All tunables in one place for CLI / UI / experiments."""
+
+    strictness: Strictness = Strictness.BALANCED
+    content_priority: ContentPriority = ContentPriority.BALANCED
+    page_type: PageType = PageType.UNKNOWN
+
+    enable_schema: bool = True
+    enable_tables: bool = True
+    enable_dvdf: bool = True
+    enable_stce: bool = True  # applied only when a SiteChromeModel is provided
+    inject_heading_paths: bool = False
+
+    # Density
+    min_chars: int = 40
+    max_link_density: float = 0.55
+    min_text_density: float = 0.12
+
+    # Chrome / size guards
+    max_noise_block_chars: int = 4000  # do not strip wrappers larger than this via class heuristics
+    max_rejected_drop_chars: int = 400  # density/DVDF rejected nodes dropped only if shorter
+
+    # DVDF
+    dvdf_threshold: float = 0.42
+
+    # STCE (site-template chrome elimination)
+    stce_min_pages: int = 5
+    stce_frequency: float = 0.80  # appear in ≥80% of site pages
+    stce_max_block_chars: int = 1500  # never STCE-drop long shared content blocks
+    stce_require_chrome_features: bool = True  # frequency alone is not enough
+
+    def density_config(self) -> DensityConfig:
+        return DensityConfig(
+            min_chars=self.min_chars,
+            max_link_density=self.max_link_density,
+            min_text_density=self.min_text_density,
+            prefer_main=True,
+        )
+
+    @classmethod
+    def from_priority(
+        cls,
+        priority: ContentPriority | str = ContentPriority.BALANCED,
+        page_type: PageType | str = PageType.UNKNOWN,
+        **overrides: object,
+    ) -> PipelineConfig:
+        """Build config from the user-facing precision/coverage flag."""
+        if isinstance(priority, str):
+            priority = ContentPriority(priority)
+        cfg = cls.from_strictness(
+            priority_to_strictness(priority),
+            page_type,
+            **overrides,
+        )
+        cfg.content_priority = priority
+        return cfg
+
+    @classmethod
+    def from_strictness(
+        cls,
+        strictness: Strictness | str = Strictness.BALANCED,
+        page_type: PageType | str = PageType.UNKNOWN,
+        **overrides: object,
+    ) -> PipelineConfig:
+        if isinstance(strictness, str):
+            strictness = Strictness(strictness)
+        if isinstance(page_type, str):
+            page_type = PageType(page_type)
+
+        # Keep content_priority aligned with strictness when callers use from_strictness
+        priority = ContentPriority.BALANCED
+        if strictness == Strictness.AGGRESSIVE:
+            priority = ContentPriority.PRECISION
+        elif strictness == Strictness.LENIENT:
+            priority = ContentPriority.COVERAGE
+
+        base = cls(strictness=strictness, content_priority=priority, page_type=page_type)
+
+        if strictness == Strictness.LENIENT:
+            base.max_link_density = 0.70
+            base.min_text_density = 0.08
+            base.min_chars = 25
+            base.dvdf_threshold = 0.55
+            base.max_rejected_drop_chars = 250
+            base.stce_frequency = 0.85
+            base.enable_dvdf = True
+        elif strictness == Strictness.AGGRESSIVE:
+            base.max_link_density = 0.40
+            base.min_text_density = 0.18
+            base.min_chars = 50
+            base.dvdf_threshold = 0.32
+            base.max_rejected_drop_chars = 600
+            base.stce_frequency = 0.55
+            base.stce_max_block_chars = 3500
+        else:  # BALANCED
+            pass
+
+        # Page-type nudges (on top of strictness)
+        if page_type == PageType.MARKETING:
+            base.max_link_density = min(base.max_link_density, 0.50)
+            base.dvdf_threshold = min(base.dvdf_threshold, 0.38)
+        elif page_type == PageType.DOCS:
+            base.max_link_density = max(base.max_link_density, 0.60)  # docs side nav is linky but keep body
+            base.enable_dvdf = True
+        elif page_type == PageType.PRICING:
+            base.enable_tables = True
+            base.max_rejected_drop_chars = min(base.max_rejected_drop_chars, 300)
+
+        for key, value in overrides.items():
+            if not hasattr(base, key):
+                raise ValueError(f"Unknown PipelineConfig field: {key}")
+            setattr(base, key, value)
+        return base
+
+
+def infer_page_type(url: str | None, html: str = "") -> PageType:
+    blob = f"{url or ''} {html[:2000]}".lower()
+    if any(k in blob for k in ("/pricing", "pricing", "on-demand", "price-list")):
+        return PageType.PRICING
+    if any(k in blob for k in ("/docs", "documentation", "userguide", "developer", "/learn/")):
+        return PageType.DOCS
+    if any(k in blob for k in ("/products", "/product/", "solutions", "platform")):
+        return PageType.MARKETING
+    if any(k in blob for k in ("/blog", "resources", "customers")):
+        return PageType.HUB
+    return PageType.UNKNOWN