dirag 0.1.0__py3-none-any.whl

dirag/__init__.py

@@ -0,0 +1,3 @@
+"""dirag: search a folder of PDF books and read the page each result comes from."""
+
+__version__ = "0.1.0"

dirag/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from .cli import main
+
+sys.exit(main())

dirag/answer.py

@@ -0,0 +1,129 @@
+"""The answer: quotes chosen by a language model, verified against the passages and located on the page.
+
+The model sees the top passages of a search and returns passage numbers with
+the words it quotes from each, plus a response of one to three sentences citing
+them by number. Nothing
+it writes is shown as a quote:
+
+  - Each quote is matched against its passage on letters and digits only, so
+    case, punctuation, spacing and line-break hyphens do not matter. A quote
+    with "..." is matched part by part, in order. A quote that does not match
+    is dropped.
+  - The text shown for a quote is the passage's own text for the matched span.
+  - The quote is located on its page through the page's words, giving one
+    rectangle per line for the reader and the page image to mark.
+  - Citations in the response to passages with no kept quote are removed, the
+    rest renumbered to the quotes. The response is kept only when at least one
+    citation remains; otherwise only the quotes are shown. With no kept quote the answer
+    is empty.
+"""
+
+import json
+import re
+
+import pymupdf
+
+from . import llm
+
+DEPTH = 10
+MAX_QUOTES = 5
+MIN_ALNUM = 15
+
+SYSTEM = ("You answer a question using only the numbered book passages given. Reply with JSON only:\n"
+          '{"quotes": [{"n": <passage number>, "text": "<words copied exactly from that passage>"}], '
+          '"summary": "<one to three sentences answering the question, each citing passages like [2]>"}\n'
+          "Copy every quote word for word from its passage; use ... only to skip words inside a quote. "
+          "Cite only passages you quoted. "
+          f"Each quote is one to three sentences. Order quotes from most to least useful; at most {MAX_QUOTES}. "
+          'If no passage answers the question, reply {"quotes": [], "summary": ""}.')
+
+
+def _alnum(text):
+    """Letters and digits of `text`, lowercased, with the index in `text` of each one."""
+    chars, where = [], []
+    for i, c in enumerate(text):
+        if c.isalnum():
+            chars.append(c.lower())
+            where.append(i)
+    return "".join(chars), where
+
+
+def locate(quote, text):
+    """(start, end) of `quote` within `text`, matched on letters and digits, or None."""
+    haystack, where = _alnum(text)
+    parts = [p for p in (_alnum(part)[0] for part in re.split(r"\.\.\.|\u2026", quote)) if p]
+    if not parts or sum(len(p) for p in parts) < MIN_ALNUM:
+        return None
+    start, at = None, 0
+    for part in parts:
+        found = haystack.find(part, at)
+        if found == -1:
+            return None
+        start = found if start is None else start
+        at = found + len(part)
+    return where[start], where[at - 1] + 1
+
+
+def page_rects(path, page_number, span_text):
+    """Line rectangles of `span_text` on a page, in PDF points, found through the page's words."""
+    doc = pymupdf.open(path)
+    try:
+        words = doc.load_page(page_number - 1).get_text("words")
+    finally:
+        doc.close()
+    joined, owner = [], []
+    for index, word in enumerate(words):
+        letters = _alnum(word[4])[0]
+        joined.append(letters)
+        owner.extend([index] * len(letters))
+    needle = _alnum(span_text)[0]
+    found = "".join(joined).find(needle)
+    if found == -1 or not needle:
+        return []
+    lines = {}
+    for index in sorted(set(owner[found:found + len(needle)])):
+        x0, y0, x1, y1, _, block, line, _ = words[index]
+        box = lines.setdefault((block, line), [x0, y0, x1, y1])
+        box[:] = [min(box[0], x0), min(box[1], y0), max(box[2], x1), max(box[3], y1)]
+    return [[round(v, 2) for v in box] for box in lines.values()]
+
+
+def compose(question, rows, file_of):
+    """{"summary", "quotes"} for `rows` (search rows, best first). `file_of(rel)` gives a book's PDF path.
+
+    Raises llm.LLMError when the model call fails.
+    """
+    top = rows[:DEPTH]
+    listing = "\n\n".join(f"[{n}] {row['book']} > {row['section'] or ''}, p. {row['page']}\n{row['text']}"
+                          for n, row in enumerate(top, 1))
+    reply = llm.chat_json(SYSTEM, f"Question: {question}\n\nPassages:\n\n{listing}")
+    quotes, numbers = [], {}
+    for item in reply.get("quotes", []) if isinstance(reply, dict) else []:
+        try:
+            n, said = int(item["n"]), str(item["text"])
+        except (KeyError, TypeError, ValueError):
+            continue
+        if not 1 <= n <= len(top) or len(quotes) >= MAX_QUOTES:
+            continue
+        row = top[n - 1]
+        span = locate(said, row["text"])
+        if span is None:
+            continue
+        text = row["text"][span[0]:span[1]]
+        if any(q["chunk_id"] == row["id"] and q["text"] == text for q in quotes):
+            continue
+        path = file_of(row["path"])
+        quotes.append({"n": len(quotes) + 1, "chunk_id": row["id"], "book_id": row["book_id"], "book": row["book"],
+                       "year": row["year"], "section": row["section"], "page": row["page"], "text": text,
+                       "rects": (page_rects(path, row["page"], text) if path else []) or _bbox(row)})
+        numbers.setdefault(n, quotes[-1]["n"])
+    summary = str(reply.get("summary") or "").strip() if quotes else ""
+    summary = re.sub(r"\s*\[(\d+)\]", lambda m: f" [{numbers[int(m.group(1))]}]" if int(m.group(1)) in numbers else "",
+                     summary).strip()
+    if not re.search(r"\[\d+\]", summary):
+        summary = ""
+    return {"summary": summary, "quotes": quotes}
+
+
+def _bbox(row):
+    return [json.loads(row["bbox_json"])] if row["bbox_json"] else []

dirag/cli.py

@@ -0,0 +1,126 @@
+"""Command line: `dirag [serve]`, `dirag index`, `dirag find`, `dirag where`, `dirag toc`."""
+
+import argparse
+import shutil
+import sys
+import textwrap
+
+from . import __version__, config
+
+
+def cmd_serve(args):
+    from .server import serve
+    if args.library:
+        config.set_library(args.library)
+    return serve(args.host, args.port, browser=not args.no_browser)
+
+
+def cmd_index(args):
+    from .index import run
+    root = config.require_library(args.library)
+    action = "reindex" if args.reindex else "rechunk" if args.rechunk else "update"
+    return run(root, config.index_path(root), action=action, limit=args.limit)
+
+
+def _index():
+    return config.index_path(config.require_library())
+
+
+def cmd_find(args):
+    from . import search
+    conn = search.connect(_index())
+    try:
+        rows = search.passages(conn, args.query, args.mode, args.rerank, limit=args.k, per_book=args.per_book)
+    finally:
+        conn.close()
+    if not rows:
+        print("No matches.")
+        return 0
+    width = min(shutil.get_terminal_size((100, 24)).columns, 100)
+    for number, row in enumerate(rows, 1):
+        print(f"{number:>3}  {row['book'][:66]}   p.{row['page']}{'   (dropped by llm)' if row.get('dropped') else ''}")
+        if row["section"]:
+            print(f"     {row['section'][:90]}")
+        print(textwrap.fill(row["text"], width=width, initial_indent="     ", subsequent_indent="     "))
+        print("")
+    return 0
+
+
+def cmd_where(args):
+    from . import search
+    conn = search.connect(_index())
+    try:
+        ranked = search.chapters(conn, search.candidates(conn, args.query, args.mode), limit=args.k)
+    finally:
+        conn.close()
+    if not ranked:
+        print("No matches.")
+        return 0
+    print(f"{'hits':>5}  {'pages':>11}  book / chapter")
+    for entry in ranked:
+        span = f"{min(entry['pages'])}-{max(entry['pages'])}"
+        stamp = f" ({entry['year']})" if entry["year"] else ""
+        print(f"{entry['hits']:>5}  {span:>11}  {entry['book'][:60]}{stamp}")
+        print(f"{'':>5}  {'':>11}  {entry['chapter'][:70]}")
+    return 0
+
+
+def cmd_toc(args):
+    import pymupdf
+    from .toc import chapters_of, sections_of
+    doc = pymupdf.open(args.pdf)
+    try:
+        sections = sections_of(doc)
+        if not sections:
+            print("No usable table of contents.")
+            return 0
+        for chapter in chapters_of(sections):
+            print(f"p.{chapter['start_page']:>5}-{chapter['end_page']:>5}  {chapter['sections']:>4} sec  {chapter['title'][:80]}")
+    finally:
+        doc.close()
+    return 0
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="dirag", description="Search a folder of PDF books and read the page.")
+    parser.add_argument("--version", action="version", version=f"dirag {__version__}")
+    sub = parser.add_subparsers(dest="command")
+
+    serve = sub.add_parser("serve", help="run the web app (the default)")
+    serve.add_argument("--library", help="the folder of PDFs; remembered once given")
+    serve.add_argument("--host", default="127.0.0.1")
+    serve.add_argument("--port", type=int, default=8008)
+    serve.add_argument("--no-browser", action="store_true", help="do not open the browser")
+    serve.set_defaults(func=cmd_serve)
+
+    index = sub.add_parser("index", help="index new and changed PDFs, drop deleted ones")
+    index.add_argument("--library", help="the folder of PDFs; remembered once given")
+    index.add_argument("--limit", type=int, default=0, help="only the first N PDFs")
+    mode = index.add_mutually_exclusive_group()
+    mode.add_argument("--reindex", action="store_true", help="parse and embed every PDF again")
+    mode.add_argument("--rechunk", action="store_true", help="rebuild passages and vectors from cached pages")
+    index.set_defaults(func=cmd_index)
+
+    find = sub.add_parser("find", help="passages for a query")
+    find.add_argument("query")
+    find.add_argument("-k", type=int, default=10)
+    find.add_argument("--per-book", type=int, default=2)
+    find.add_argument("--mode", choices=("lexical", "semantic", "hybrid"), default="hybrid")
+    find.add_argument("--rerank", choices=("off", "neural", "llm"), default="off")
+    find.set_defaults(func=cmd_find)
+
+    where = sub.add_parser("where", help="chapters for a query")
+    where.add_argument("query")
+    where.add_argument("-k", type=int, default=10)
+    where.add_argument("--mode", choices=("lexical", "semantic", "hybrid"), default="hybrid")
+    where.set_defaults(func=cmd_where)
+
+    toc = sub.add_parser("toc", help="the chapter map of one PDF")
+    toc.add_argument("pdf")
+    toc.set_defaults(func=cmd_toc)
+
+    argv = list(sys.argv[1:] if argv is None else argv)
+    if not argv or (argv[0] not in sub.choices and argv[0] not in ("-h", "--help", "--version")):
+        argv = ["serve", *argv]
+    args = parser.parse_args(argv)
+    return args.func(args)

dirag/config.py

@@ -0,0 +1,85 @@
+"""Where dirag keeps its files, and the one library folder it serves.
+
+Two directories, both overridable:
+
+    DIRAG_HOME    derived data: one index per library folder (indexes/), the
+                  embedding model cache (models/) and the indexing job record
+                  (job.json). Default ~/.local/share/dirag. Rebuildable.
+    DIRAG_STATE   user state: config.json, positions.json, bookmarks.json,
+                  cards.json. Default DIRAG_HOME. Small, textual, not derived.
+
+The library is a single folder of PDFs. It is resolved in this order: the
+--library option, the DIRAG_LIBRARY variable, then config.json. Passing
+--library, or choosing a folder in the app, saves it to config.json. With
+DIRAG_LIBRARY set the folder is fixed and the app cannot change it.
+
+The app's folder picker lists folders under DIRAG_BROWSE_ROOT only (default:
+the home directory).
+"""
+
+import hashlib
+import json
+import os
+import sys
+from pathlib import Path
+
+HOME = Path(os.getenv("DIRAG_HOME") or Path.home() / ".local" / "share" / "dirag").expanduser()
+STATE = Path(os.getenv("DIRAG_STATE") or HOME).expanduser()
+MODELS = HOME / "models"
+JOB = HOME / "job.json"
+CONFIG = STATE / "config.json"
+FIXED = bool(os.getenv("DIRAG_LIBRARY"))
+BROWSE_ROOT = Path(os.getenv("DIRAG_BROWSE_ROOT") or Path.home()).expanduser().resolve()
+
+
+def read_json(path, empty):
+    """The file's content, or `empty` when it is missing, unreadable or the wrong type."""
+    try:
+        data = json.loads(Path(path).read_text(encoding="utf-8"))
+        return data if isinstance(data, type(empty)) else empty
+    except FileNotFoundError:
+        return empty
+    except (OSError, ValueError) as exc:
+        print(f"{Path(path).name} unusable ({exc}); ignoring it", file=sys.stderr)
+        return empty
+
+
+def write_json(path, data):
+    """Write through a temp file and rename, so a reader never sees half a file."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    tmp.write_text(json.dumps(data, indent=1, sort_keys=True) + "\n", encoding="utf-8")
+    tmp.replace(path)
+
+
+def index_path(root):
+    """The index file for a library folder: its name plus a hash of its absolute path."""
+    digest = hashlib.sha1(str(root).encode()).hexdigest()[:10]
+    return HOME / "indexes" / f"{root.name or 'root'}-{digest}.sqlite3"
+
+
+def set_library(path):
+    """Save `path` as the library folder and return it as an absolute Path."""
+    root = Path(path).expanduser().resolve()
+    if not root.is_dir():
+        raise SystemExit(f"Not a folder: {root}")
+    settings = read_json(CONFIG, {})
+    settings["library"] = str(root)
+    write_json(CONFIG, settings)
+    return root
+
+
+def library(override=None):
+    """The library folder as an absolute Path, or None when none is set."""
+    if override:
+        return set_library(override)
+    value = os.getenv("DIRAG_LIBRARY") or read_json(CONFIG, {}).get("library")
+    return Path(value).expanduser().resolve() if value else None
+
+
+def require_library(override=None):
+    root = library(override)
+    if root is None:
+        raise SystemExit("No library folder set. Pass --library PATH once; it is remembered.")
+    return root

dirag/embed.py

@@ -0,0 +1,130 @@
+"""The local models (the embedder here, the cross-encoder in rerank.py) and where they run.
+
+They run on an NVIDIA GPU when one is present and working, else on the CPU.
+Present means an NVIDIA driver is loaded and onnxruntime has its CUDA provider;
+working means a first inference succeeds. A GPU that fails that check is
+reported once and the CPU is used instead.
+
+    DIRAG_DEVICE=cpu        never use the GPU
+    DIRAG_EMBED_THREADS     CPU threads for onnxruntime; default all cores
+"""
+
+import os
+import shutil
+import subprocess
+import sys
+import threading
+import warnings
+from pathlib import Path
+
+from . import config
+
+MODEL = "BAAI/bge-small-en-v1.5"
+DIM = 384
+THREADS = int(os.getenv("DIRAG_EMBED_THREADS")) if os.getenv("DIRAG_EMBED_THREADS") else None
+
+_model = None
+_lock = threading.Lock()
+_gpu = None                       # None until decided, then True or False
+
+
+def gpu_name():
+    """The NVIDIA GPU's name, or None when there is no usable one. Does not load onnxruntime.
+
+    nvidia-smi naming a GPU is the test; without nvidia-smi, the GPU device file is.
+    Driver files alone (as in a container without the GPU passed through) do not count.
+    """
+    if os.getenv("DIRAG_DEVICE", "").lower() == "cpu":
+        return None
+    if shutil.which("nvidia-smi"):
+        try:
+            out = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
+                                 capture_output=True, text=True, timeout=5)
+            names = out.stdout.strip().splitlines()
+            return names[0] if out.returncode == 0 and names else None
+        except (OSError, subprocess.SubprocessError):
+            return None
+    return "NVIDIA GPU" if Path("/dev/nvidia0").exists() else None
+
+
+def use_gpu():
+    """Whether the local models should try the GPU."""
+    global _gpu
+    if _gpu is None:
+        _gpu = False
+        if gpu_name():
+            import onnxruntime
+            if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
+                # CUDA and cuDNN installed as pip packages load from here.
+                try:
+                    onnxruntime.preload_dlls()
+                except Exception:
+                    pass
+                _gpu = True
+    return _gpu
+
+
+def build(make, probe):
+    """make(**options) on the GPU when it works, else on the CPU. probe(model) runs one inference.
+
+    onnxruntime falls back to the CPU on its own when the CUDA provider cannot
+    start, with only a warning; that warning counts as the GPU not working.
+    """
+    global _gpu
+    import onnxruntime
+    # dirag reports the device itself; onnxruntime's own provider errors are noise here.
+    onnxruntime.set_default_logger_severity(4)
+    if use_gpu():
+        try:
+            with warnings.catch_warnings(record=True) as caught:
+                warnings.simplefilter("always")
+                model = make(cuda=True, device_ids=[0])
+                probe(model)
+            failed = [str(w.message) for w in caught if "CUDAExecutionProvider" in str(w.message)]
+            if not failed:
+                return model
+            reason = failed[0]
+        except Exception as exc:
+            reason = str(exc)
+        print(f"dirag: the GPU did not work ({reason.splitlines()[0][:160]}); using the CPU", file=sys.stderr, flush=True)
+        _gpu = False
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", RuntimeWarning)
+        model = make(threads=THREADS, providers=["CPUExecutionProvider"])
+        probe(model)
+    return model
+
+
+def _load():
+    global _model
+    with _lock:
+        if _model is None:
+            from fastembed import TextEmbedding
+            _model = build(lambda **o: TextEmbedding(model_name=MODEL, cache_dir=str(config.MODELS), **o),
+                           lambda m: list(m.embed(["probe"])))
+    return _model
+
+
+def batch_size():
+    """Passages per embedding call: large on the GPU, small on the CPU so a stop is honoured within seconds."""
+    _load()
+    return 256 if _gpu else 32
+
+
+def passages(texts):
+    """Vectors for a batch of passages, as lists of floats."""
+    texts = list(texts)
+    return [vector.tolist() for vector in _load().embed(texts)] if texts else []
+
+
+def query(text):
+    """Vector for one query. bge models prefix queries with an instruction; query_embed adds it."""
+    return list(_load().query_embed([text]))[0].tolist()
+
+
+def preflight():
+    """Load the embedder before a long run, so a broken setup fails before the first book."""
+    try:
+        _load()
+    except Exception as exc:
+        raise SystemExit(f"Embedding is not working:\n  {exc}")