130 lines
5.1 KiB
Python
130 lines
5.1 KiB
Python
"""Framework-agnostic local-folder retriever — the D7-portable data seam.
|
|
|
|
**NO** ``agent_framework``, **NO** ``mcp`` — pure stdlib. Documents are chunked at ingest
|
|
with EXACT character spans, so every retrieved chunk carries a locator that slices its own
|
|
snippet byte-for-byte (citations exact by construction — research 02). Scoring is
|
|
keyword/substring overlap (no heavy vector deps — D5/D6).
|
|
|
|
Path-security (this module reads a docs folder; the EscapeRoute CVE class, research 02
|
|
Dim 5): every accessed path is canonicalised with ``os.path.realpath`` and boundary-checked
|
|
against the canonicalised ``docs_dir``. A ``..`` traversal, a symlink escaping the folder,
|
|
and a sibling directory sharing a name prefix (``/a/docs`` vs ``/a/docs-evil``) are all
|
|
rejected fail-closed. Native file APIs only.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
_CHUNK_CHARS = 240 # paragraph-ish window; small + deterministic for the synthetic corpus
|
|
|
|
|
|
class PathSecurityError(RuntimeError):
|
|
"""A requested path escapes ``docs_dir`` (traversal / symlink / prefix-collision)."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TextSpan:
|
|
"""An exact half-open character span ``[start_index, end_index)`` into a document's
|
|
text. The SINGLE owner of this type — ``provenance.Citation.locator`` imports it from
|
|
here (Step 6)."""
|
|
|
|
start_index: int
|
|
end_index: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetrievedChunk:
|
|
"""One retrieved chunk with an exact-by-construction citation locator."""
|
|
|
|
file: str
|
|
locator: TextSpan
|
|
snippet: str
|
|
score: float
|
|
|
|
|
|
def is_within_dir(candidate: str, docs_dir: str) -> bool:
|
|
"""True iff the realpath of ``candidate`` is inside the realpath of ``docs_dir``.
|
|
|
|
Fail-closed against symlink-escape and prefix-collision: uses ``os.path.commonpath`` on
|
|
canonical paths, so a sibling dir sharing a name prefix (``/a/docs-evil`` vs ``/a/docs``)
|
|
is correctly rejected (a naive ``startswith`` would not)."""
|
|
real_root = os.path.realpath(docs_dir)
|
|
real_candidate = os.path.realpath(candidate)
|
|
try:
|
|
return os.path.commonpath([real_root, real_candidate]) == real_root
|
|
except ValueError:
|
|
# Different drives / mixed absolute-relative -> not within.
|
|
return False
|
|
|
|
|
|
def safe_resolve(docs_dir: str, relative: str) -> str:
|
|
"""Resolve ``relative`` against ``docs_dir`` and return its canonical path, or raise
|
|
``PathSecurityError`` if it escapes the folder (``..`` traversal, symlink escape, or a
|
|
prefix-collision sibling). Fail-closed."""
|
|
target = os.path.join(os.path.realpath(docs_dir), relative)
|
|
if not is_within_dir(target, docs_dir):
|
|
raise PathSecurityError(f"path escapes docs_dir: {relative!r}")
|
|
return os.path.realpath(target)
|
|
|
|
|
|
def _resolve_docs_dir(docs_dir: str) -> str:
|
|
real = os.path.realpath(docs_dir)
|
|
if not os.path.isdir(real):
|
|
raise ValueError(f"docs_dir is not a directory: {docs_dir}")
|
|
return real
|
|
|
|
|
|
def _chunk_text(text: str) -> list[TextSpan]:
|
|
"""Split text into fixed-size character windows, each an exact half-open span."""
|
|
spans: list[TextSpan] = []
|
|
for start in range(0, len(text), _CHUNK_CHARS):
|
|
end = min(start + _CHUNK_CHARS, len(text))
|
|
if end > start:
|
|
spans.append(TextSpan(start_index=start, end_index=end))
|
|
return spans
|
|
|
|
|
|
def _score(query: str, snippet: str) -> float:
|
|
"""Keyword overlap: fraction of distinct query tokens present in the snippet (lowered).
|
|
Deterministic, no vectors."""
|
|
q_tokens = {t for t in query.lower().split() if t}
|
|
if not q_tokens:
|
|
return 0.0
|
|
blob = snippet.lower()
|
|
hits = sum(1 for t in q_tokens if t in blob)
|
|
return hits / len(q_tokens)
|
|
|
|
|
|
def retrieve(query: str, docs_dir: str, top_k: int = 3) -> list[RetrievedChunk]:
|
|
"""Return the top-``top_k`` chunks across the docs folder ranked by keyword overlap.
|
|
|
|
Every file is boundary-checked against ``docs_dir`` before it is read (fail-closed: a
|
|
symlink whose realpath escapes the folder is skipped, never opened). Locators are exact
|
|
by construction: ``snippet == file_text[start:end]``. Deterministic: ties break by
|
|
``(file, start_index)``."""
|
|
if top_k <= 0:
|
|
raise ValueError(f"top_k must be positive, got {top_k}")
|
|
root = _resolve_docs_dir(docs_dir)
|
|
candidates: list[RetrievedChunk] = []
|
|
for entry in sorted(Path(root).rglob("*")):
|
|
if not entry.is_file():
|
|
continue
|
|
if not is_within_dir(str(entry), root):
|
|
continue # symlink escaping the folder -> fail closed, never read
|
|
try:
|
|
text = entry.read_text(encoding="utf-8")
|
|
except UnicodeDecodeError:
|
|
continue # skip binary
|
|
rel = os.path.relpath(os.path.realpath(str(entry)), root)
|
|
for span in _chunk_text(text):
|
|
snippet = text[span.start_index : span.end_index]
|
|
candidates.append(
|
|
RetrievedChunk(
|
|
file=rel, locator=span, snippet=snippet, score=_score(query, snippet)
|
|
)
|
|
)
|
|
ranked = sorted(candidates, key=lambda c: (-c.score, c.file, c.locator.start_index))
|
|
return ranked[:top_k]
|