feat(index,serve): department-filtered search for local document indexes (#115)

- pixelrag index build records a "department" field in articles.json,
  derived from the first sub-directory under the local source root
  (docs/<department>/<file> layout)
- serve: new optional "department" field on POST /search pre-filters
  inside FAISS via IDSelectorBatch (SearchParametersIVF on IVF indexes),
  so n_docs results are guaranteed without over-fetching; unknown
  department returns 404 with the list of valid names; indexes built
  without department metadata return a clear 400
- new GET /departments endpoint listing departments with document counts
- fix: rename PDF tile dirs to the pipeline position index. render_pdf
  names tile dirs after the PDF filename stem, and non-numeric stems made
  embed fall back to salted hash() article ids, breaking the positional
  articles.json mapping for every local PDF source
- fix: serve now resolves PDF page images (tile_XXXX.jpg) when
  materialized chunk files do not exist, and /tile returns the correct
  MIME type for JPEG tiles
- tests: build-side department mapping and serve-side selector filtering
  (flat + IVF), 404/400 error paths

Co-authored-by: Andy Lee <andylizf@outlook.com>

Author: locphamnguyen

index/src/pixelrag_index/pipelines.py

@@ -11,6 +11,28 @@
 logger = logging.getLogger("pixelrag-index")
 
 
+def _department_of(article: dict, source_root: str) -> str:
+    """Department = first sub-directory under the source root holding the file.
+
+    Files directly under the root, non-local documents (web URLs), or paths
+    outside the root have no department ("").
+    """
+    raw = article.get("path") or ""
+    if not raw:
+        url = article.get("url") or ""
+        if url.startswith("file://"):
+            from urllib.parse import unquote, urlparse
+
+            raw = unquote(urlparse(url).path)
+    if not raw or not source_root:
+        return ""
+    try:
+        rel = Path(raw).resolve().relative_to(Path(source_root).expanduser().resolve())
+    except ValueError:
+        return ""
+    return rel.parts[0] if len(rel.parts) > 1 else ""
+
+
 def build(config: dict, limit: int | None = None, force: bool = False) -> Path:
     """Build a searchable FAISS index from a document source.
 
@@ -174,6 +196,9 @@ def _repl(m: re.Match) -> str:
     # Render PDFs — use idx as tile directory name (like URLs) so directory
     # names are always the numeric article_id.
     for idx, doc in pdf_docs:
+        out_dir = tiles_dir / f"{idx}.png.tiles"
+        if (out_dir / "tiles.json").exists():
+            continue  # already rendered on a previous run
         try:
             render_pdf(doc.path, str(tiles_dir), stem=str(idx))
         except Exception as e:
@@ -238,6 +263,7 @@ def _repl(m: re.Match) -> str:
     # as doc IDs, which are not numeric. int() on a filename stem raises ValueError
     # and crashes the entire index build step.
     articles_path = output / "articles.json"
+    source_root = str(config.get("source", {}).get("path", "") or "")
     article_entries = []
     for enum_idx, a in enumerate(articles):
         title = a.get("metadata", {}).get("title", "")
@@ -247,7 +273,11 @@ def _repl(m: re.Match) -> str:
             # Fall back to original doc id (e.g. filename stem) as display title
             title = a.get("id", str(enum_idx))
         url = a.get("url", "") or a.get("path", "")
-        article_entries.append({"title": title, "url": url})
+        # Department (from the directory layout of local sources) enables
+        # server-side filtered search; empty string means "no department".
+        article_entries.append(
+            {"title": title, "url": url, "department": _department_of(a, source_root)}
+        )
     with open(articles_path, "w") as f:
         json.dump(article_entries, f)
     logger.info(

serve/src/pixelrag_serve/api.py

@@ -148,6 +148,11 @@ class SearchRequest(BaseModel):
     instruction: str | None = None  # override query embedding instruction
     include_images: bool = False  # return base64-encoded tile images
     articles_only: bool = False  # drop Wikipedia meta pages (Portal:, List_of_, …)
+    # Restrict search to one department (articles.json "department" field, set by
+    # `pixelrag index build` from the source directory layout). Pre-filters inside
+    # FAISS via an IDSelector — not a post-filter, so n_docs results are guaranteed
+    # when the department has enough tiles.
+    department: str | None = None
 
 
 # Wikipedia meta/aggregator pages that pollute "find the article" results.
@@ -166,6 +171,38 @@ def _is_meta(url: str) -> bool:
     return bool(_META_RE.search(url))
 
 
+def _department_positions(department: str) -> np.ndarray:
+    """Vector positions (FAISS row ids) belonging to a department, cached per dept."""
+    dept_to_aids = _state.get("dept_to_aids") or {}
+    if not dept_to_aids:
+        raise HTTPException(
+            status_code=400,
+            detail="Index was built without department metadata; rebuild with "
+            "`pixelrag index build` from a directory-per-department source.",
+        )
+    cache = _state.setdefault("dept_positions", {})
+    if department not in cache:
+        aids = dept_to_aids.get(department)
+        if aids is None:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Unknown department {department!r}. "
+                f"Available: {sorted(dept_to_aids)}",
+            )
+        article_ids = _state["metadata"]["article_ids"]
+        cache[department] = np.where(np.isin(article_ids, aids))[0].astype("int64")
+    return cache[department]
+
+
+def _department_search_params(department: str, nprobe: int) -> "faiss.SearchParameters":
+    """Build FAISS search params that pre-filter to one department's vectors."""
+    sel = faiss.IDSelectorBatch(_department_positions(department))
+    index = _state["index"]
+    if isinstance(index, faiss.IndexIVF):
+        return faiss.SearchParametersIVF(sel=sel, nprobe=nprobe)
+    return faiss.SearchParameters(sel=sel)
+
+
 class Hit(BaseModel):
     score: float
     vector_id: int
@@ -340,6 +377,12 @@ def _resolve_path(article_id: int, tile_index: int, chunk_index: int) -> str:
     if os.path.exists(flat_path):
         return flat_path
 
+    # PDF pipeline stores whole pages as tile_XXXX.jpg (one chunk per page,
+    # no materialized chunk files) — fall back to the page image.
+    page_path = os.path.join(tiles_dir, tiles_dirname, f"tile_{tile_index:04d}.jpg")
+    if os.path.exists(page_path):
+        return page_path
+
     # Sharded layout: tiles_dir/shard_XXX/sub/{article_id}.png.tiles/chunk_XXXX_YY.png
     top_shard = article_id // shard_size
     top_shard_dir = os.path.join(tiles_dir, f"shard_{top_shard:03d}")
@@ -426,7 +469,12 @@ async def search(req: SearchRequest):
         fetch_k = req.n_docs * 5
     else:
         fetch_k = req.n_docs
-    distances, indices = index.search(query_vectors, fetch_k)
+    if req.department:
+        # Department pre-filter: FAISS only scores vectors of that department.
+        params = _department_search_params(req.department, index.nprobe)
+        distances, indices = index.search(query_vectors, fetch_k, params=params)
+    else:
+        distances, indices = index.search(query_vectors, fetch_k)
 
     if req.nprobe is not None:
         index.nprobe = default_nprobe
@@ -528,6 +576,18 @@ async def health():
     return {"status": "ok"}
 
 
+@app.get("/departments")
+async def departments():
+    """Departments available for the `department` search filter, with doc counts."""
+    dept_to_aids = _state.get("dept_to_aids") or {}
+    return {
+        "departments": [
+            {"name": d, "n_documents": len(aids)}
+            for d, aids in sorted(dept_to_aids.items())
+        ]
+    }
+
+
 class ReconstructRequest(BaseModel):
     vector_ids: list[int]
 
@@ -564,7 +624,8 @@ async def tile_by_id(article_id: int, tile_index: int, chunk_index: int):
     path = _resolve_path(article_id, tile_index, chunk_index)
     if not os.path.isfile(path):
         raise HTTPException(status_code=404, detail="Tile not found")
-    return FileResponse(path, media_type="image/png")
+    media_type = "image/jpeg" if path.endswith((".jpg", ".jpeg")) else "image/png"
+    return FileResponse(path, media_type=media_type)
 
 
 # ---------------------------------------------------------------------------
@@ -612,6 +673,20 @@ def load(args):
         articles = json.load(f)
     logger.info("Loaded %d article slugs", len(articles))
 
+    # Department → article ids, for the `department` search filter. Older
+    # indexes (or web/kiwix sources) have no "department" key — the map stays
+    # empty and filtered requests get a clear 400.
+    dept_to_aids: dict[str, list[int]] = {}
+    for aid, a in enumerate(articles):
+        dept = a.get("department", "") if isinstance(a, dict) else ""
+        if dept:
+            dept_to_aids.setdefault(dept, []).append(aid)
+    if dept_to_aids:
+        logger.info(
+            "Departments: %s",
+            ", ".join(f"{d}({len(v)})" for d, v in sorted(dept_to_aids.items())),
+        )
+
     # Load embedding model
     from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
 
@@ -646,6 +721,8 @@ def load(args):
             "index": index,
             "metadata": meta,
             "articles": articles,
+            "dept_to_aids": {d: np.asarray(v) for d, v in dept_to_aids.items()},
+            "dept_positions": {},
             "processor": processor,
             "model": model,
             "device": device,

tests/test_department_filter.py

@@ -0,0 +1,130 @@
+"""Department metadata (articles.json) + FAISS IDSelector pre-filter."""
+
+import pytest
+
+pipelines = pytest.importorskip("pixelrag_index.pipelines")
+
+
+# ---------------------------------------------------------------------------
+# Build side: _department_of — directory layout → department name
+# ---------------------------------------------------------------------------
+
+
+def test_department_from_subdirectory(tmp_path):
+    (tmp_path / "hr").mkdir()
+    f = tmp_path / "hr" / "sop_tuyen_dung.pdf"
+    f.write_bytes(b"x")
+    assert pipelines._department_of({"path": str(f)}, str(tmp_path)) == "hr"
+
+
+def test_department_from_nested_subdirectory(tmp_path):
+    d = tmp_path / "ketoan" / "2026"
+    d.mkdir(parents=True)
+    f = d / "sop_thanh_toan.pdf"
+    f.write_bytes(b"x")
+    assert pipelines._department_of({"path": str(f)}, str(tmp_path)) == "ketoan"
+
+
+def test_file_at_source_root_has_no_department(tmp_path):
+    f = tmp_path / "sop.pdf"
+    f.write_bytes(b"x")
+    assert pipelines._department_of({"path": str(f)}, str(tmp_path)) == ""
+
+
+def test_department_from_file_url(tmp_path):
+    (tmp_path / "it").mkdir()
+    f = tmp_path / "it" / "huong dan.html"
+    f.write_bytes(b"x")
+    art = {"url": f.resolve().as_uri()}  # file:// with percent-encoded space
+    assert pipelines._department_of(art, str(tmp_path)) == "it"
+
+
+def test_web_url_outside_root_and_empty_root(tmp_path):
+    assert pipelines._department_of({"url": "https://ex.am/ple"}, str(tmp_path)) == ""
+    assert pipelines._department_of({"path": "/elsewhere/f.pdf"}, str(tmp_path)) == ""
+    assert pipelines._department_of({"path": "/a/b.pdf"}, "") == ""
+
+
+# ---------------------------------------------------------------------------
+# Serve side: IDSelector filter
+# ---------------------------------------------------------------------------
+
+faiss = pytest.importorskip("faiss")
+np = pytest.importorskip("numpy")
+api = pytest.importorskip("pixelrag_serve.api")
+
+DIM = 8
+# 3 articles × 2 vectors; article 0, 2 → hr; article 1 → ketoan
+ARTICLE_IDS = np.array([0, 0, 1, 1, 2, 2], dtype=np.int64)
+HR_ROWS = {0, 1, 4, 5}
+
+
+def _vectors() -> np.ndarray:
+    rng = np.random.default_rng(7)
+    v = rng.standard_normal((6, DIM)).astype(np.float32)
+    return v / np.linalg.norm(v, axis=1, keepdims=True)
+
+
+def _setup_state(index):
+    api._state.clear()
+    api._state.update(
+        {
+            "index": index,
+            "metadata": {"article_ids": ARTICLE_IDS},
+            "dept_to_aids": {"hr": np.array([0, 2]), "ketoan": np.array([1])},
+            "dept_positions": {},
+        }
+    )
+
+
+def test_department_positions_maps_articles_to_rows():
+    index = faiss.IndexFlatIP(DIM)
+    index.add(_vectors())
+    _setup_state(index)
+    assert set(api._department_positions("hr").tolist()) == HR_ROWS
+    assert set(api._department_positions("ketoan").tolist()) == {2, 3}
+
+
+def test_flat_index_filter_restricts_hits():
+    vecs = _vectors()
+    index = faiss.IndexFlatIP(DIM)
+    index.add(vecs)
+    _setup_state(index)
+    params = api._department_search_params("hr", nprobe=1)
+    _, indices = index.search(vecs[:1], 6, params=params)
+    hits = {int(i) for i in indices[0] if i != -1}
+    assert hits == HR_ROWS  # every hr row returned, nothing else
+
+
+def test_ivf_index_filter_restricts_hits():
+    vecs = _vectors()
+    quantizer = faiss.IndexFlatIP(DIM)
+    index = faiss.IndexIVFFlat(quantizer, DIM, 1, faiss.METRIC_INNER_PRODUCT)
+    index.train(vecs)
+    index.add(vecs)
+    index.nprobe = 1
+    _setup_state(index)
+    params = api._department_search_params("ketoan", nprobe=1)
+    _, indices = index.search(vecs[2:3], 6, params=params)
+    hits = {int(i) for i in indices[0] if i != -1}
+    assert hits == {2, 3}
+
+
+def test_unknown_department_raises_404():
+    index = faiss.IndexFlatIP(DIM)
+    index.add(_vectors())
+    _setup_state(index)
+    with pytest.raises(api.HTTPException) as exc:
+        api._department_positions("phong-khong-ton-tai")
+    assert exc.value.status_code == 404
+    assert "hr" in exc.value.detail  # lists available departments
+
+
+def test_index_without_department_metadata_raises_400():
+    index = faiss.IndexFlatIP(DIM)
+    index.add(_vectors())
+    _setup_state(index)
+    api._state["dept_to_aids"] = {}
+    with pytest.raises(api.HTTPException) as exc:
+        api._department_positions("hr")
+    assert exc.value.status_code == 400