diff --git a/README.md b/README.md new file mode 100644 index 0000000..fbaf678 --- /dev/null +++ b/README.md @@ -0,0 +1,360 @@ +# md-embed (c) 2024 web3dguy + +A Python script for processing Markdown files, generating embeddings, and storing them in a vector store. This tool allows you to clean, split, and embed Markdown documents using various methods and embedding models. +Features + + Data Cleaning: Removes duplicates and filters out unwanted content like '404' pages and lines containing the '©' symbol. + Flexible Input: Supports input from JSON files containing URLs and Markdown data, folders of Markdown files, or single Markdown files. + Document Splitting: Splits documents using Markdown headers or recursive character splitting. + Embedding Options: Supports embedding using HuggingFace or Ollama embeddings. + Vector Store Integration: Stores embeddings in a Chroma vector store for efficient retrieval and analysis. + Customizable Filters: Option to disable filters that remove specific content. + Logging: Generates logs for duplicates and removed files for better traceability. + + + +Installation +Prerequisites +```bash + Python 3.7 or higher + pip + Git (optional, for cloning the repository) +``` +Clone the Repository + +```bash + +git clone https://github.com/GATERAGE/mdmbed.git +cd mdmbed +``` +Install Required Packages + +Install the required Python packages using pip: + +```bash + +pip install -r requirements.txt +``` +Note: The requirements.txt file should list all the dependencies, such as tqdm, langchain, chromadb, huggingface, etc. +Usage + +Run the script using Python: + +```bash + +python md-embed.py [--filters-off] +``` +# Command-Line Arguments + + --filters-off: Disable filters that remove lines containing '©' and skip files containing both '404' and 'page not found'. + +Upon running the script, you will be prompted to choose an input method: + + JSON Input File Containing URLs and Markdown Data + Folder of Markdown Files + Single Markdown File + +JSON Input File + +If you choose Option 1, you will be asked to provide: + + Path of the JSON input file: The file should be a JSON array of objects, each containing url and markdown keys. + Path of the output folder: The folder where cleaned Markdown files and logs will be saved. + +The script will: + + Clean the data by removing duplicates. + Save the cleaned Markdown files to the specified output folder. + Generate a file_to_url.json mapping file. + Display a summary of the processing. + +Folder of Markdown Files + +If you choose Option 2, you will be asked to provide: + + Path of the folder containing Markdown files. + +The script will: + + Load all .md files from the specified folder. + Optionally filter out unwanted content. + Proceed to document splitting. + +Single Markdown File + +If you choose Option 3, you will be asked to provide: + + Path of the Markdown file. + +The script will: + + Load the specified Markdown file. + Optionally filter out unwanted content. + Proceed to document splitting. + +Document Splitting + +After loading the documents, you will be prompted to split them: + + Split Method: Choose between markdown or recursive splitting. + Remove Links: Optionally remove links from the Markdown content. + Language: Specify the programming language or language of the content. + Additional Settings: + For Markdown Splitting: + Header Levels: Specify which header levels (#, ##, etc.) to split on. + For Recursive Splitting: + Chunk Size: Specify the maximum size of each chunk. + Chunk Overlap: Specify the number of overlapping characters between chunks. + +You will have the option to preview the split data before proceeding. +Embedding and Saving + +After splitting, you will be prompted to embed and save the documents: + + Embedding Method: Choose between huggingface or ollama. + HuggingFace: Enter the embedding model name (default: all-MiniLM-L6-v2). + Ollama: Enter the Ollama model name (default: nomic-embed-text). + Persist Directory: Specify the directory to save the vector store database. + Collection Name: Enter a name for the Chroma collection. + +The script will: + + Embed the documents using the chosen embedding method. + Save the embeddings to a Chroma vector store. + Display information about the saved collections. + +Examples +Example 1: Process JSON Input File + +```bash + +python md-embed.py +``` +Choose Input Method: 1 + + Enter the path of the JSON input file: ./data/input.json + Enter the path of the output folder: ./output + +Proceed through the prompts to clean data, split documents, and embed them. +Example 2: Process Folder of Markdown Files with Filters Off + +```bash + +python md-embed.py --filters-off +``` +Choose Input Method: 2 + + Enter the path of the folder containing markdown files: ./markdown_files + +Proceed through the prompts to load, split, and embed the documents. +Contributing + +Contributions are welcome! Please follow these steps: + + Fork the repository. + + Create a new branch: + +```bash + +git checkout -b feature/your-feature-name +``` +Make your changes and commit them: + +```bash +git commit -m "Add your message" +``` +Push to the branch: + +```bash +git push origin feature/your-feature-name +``` + Open a Pull Request. + +Please make sure your code adheres to the existing style and that all tests pass. +License + +This project is licensed under the MIT License. +Acknowledgments + web3dguy + LangChain for text splitting and document handling. + HuggingFace for embedding models. + Chroma for the vector store. + TQDM for progress bars. + The open-source community for continuous support and contributions. + +# Markdown Processor and Embedder + +md-embed processes markdown files, cleans and prepares the data, splits the text into manageable chunks, and creates embeddings for use in vector databases (specifically ChromaDB). It supports multiple input methods and provides options for customizing the splitting and embedding process.
+ +## Features + +* **Multiple Input Methods:**
+ * JSON file containing URLs and markdown data
+ * Folder of markdown files
+ * Single markdown file
+* **Data Cleaning:**
+ * Removes duplicate entries based on URL section titles
+ * Handles encoding issues
+ * Sanitizes filenames for safe saving
+ * Optionally filters out files containing "404" and "page not found" (can be disabled)
+ * Removes lines containing the copyright symbol "©"
+* **Text Splitting:**
+ * **Markdown Header Splitting:** Splits text based on specified markdown header levels (e.g., `#`, `##`). Allows for custom header level selection. Preserves header hierarchy in metadata
+ * **Recursive Character Text Splitting:** Splits text into chunks of specified size and overlap
+ * **Link Removal:** Optionally removes markdown links, keeping only the link text
+* **Embedding Generation:*
+ * Supports **Hugging Face** embeddings (using `langchain_huggingface`). Defaults to `all-MiniLM-L6-v2`
+ * Supports **Ollama** embeddings (using `langchain_community`). Defaults to `nomic-embed-text`, requires a local Ollama server running at `http://localhost:11434`
+* **Vector Database Integration:**
+ * Uses **ChromaDB** (`langchain_chroma`) to store embeddings and associated metadata
+ * Allows specifying the collection name and persistence directory
+ * Handles large datasets by processing in batches
+* **Logging:**
+ * Comprehensive logging through the `logging` module
+* **Duplicate Logs**:
+ * Writes URLs with duplicate sections to a log
+* **Removed Files Logs**
+ * Write to a log files that have been removed due to filters
+ +## Requirements + +* Python 3.7+
+* `langchain` (various components - see import statements)
+* `chromadb`
+* `tqdm`
+* `beautifulsoup4` (if you were scraping, but this script doesn't actually use it)
+* `requests` (if you were scraping, but this script doesn't actually use it)
+ +To install the required packages, run:
+ +```bash +pip install langchain langchain-chroma langchain-huggingface tqdm +``` +```markdown +If you are planning to use Ollama, you need to: +Install Ollama by following the instructions provided at Ollama's official website. +Run an Ollama server locally on port 11434 +``` +md-embed can be run from the command line. It provides a command-line interface using argparse with the following option:
+--filters-off: Disables the "404" and "©" filters
+The script will then guide you through a series of interactive prompts to configure the processing:
+Input Method Selection: Choose between JSON input, a folder of markdown files, or a single markdown file
+Input File/Folder/URL: Provide the path to the input file or folder, as appropriate
+Output Folder (for JSON input): Specify the directory where cleaned markdown files will be saved
+Data Cleaning Options: The script will show total entires and total duplicates
+Language: Specify the primary language of the input files (e.g., "TypeScript", "Python")
+Splitting Method: Choose between "markdown" (header-based splitting) and "recursive" (chunk size and overlap)
+Markdown Splitting Options (if applicable):
+Remove Links: Choose whether to remove markdown links
+Header Levels: Specify which header levels to split on (e.g., "1,2,3" for #, ##, and ###). Enter "all" for all header levels
+Recursive Splitting Options (if applicable):
+Remove Links: Choose whether to remove markdown links
+Chunk Size: Specify the desired chunk size (in characters)
+Chunk Overlap: Specify the desired chunk overlap (in characters)
+Preview Splits: Choose whether to preview the split data ("yes", "full", or "no")
+Split Again: You'll be prompted to continue or modify the settings
+Embedding Method: Choose between "huggingface" and "ollama"
+Embedding Model (Hugging Face): Enter the Hugging Face model name (defaults to all-MiniLM-L6-v2)
+Embedding Model (Ollama): Enter the Ollama model name (defaults to nomic-embed-text)
+Persistence Directory: Specify the directory where the ChromaDB database will be stored
+Collection Name: Choose a name for the ChromaDB collection
+Example (JSON Input):
+```bash +python md-embed.py +``` +Follow the prompts, providing the necessary information (input file, output folder, embedding choices, etc.)
+Example (Disabling Filters): +```bash +python md-embed.py --filters-off +``` +Cleaned Markdown Files (JSON Input): If using JSON input, the script will save cleaned markdown files to the specified output folder
+ChromaDB Database: The script will create a ChromaDB database in the specified persistence directory, containing the embeddings and metadata
+Logs: The logs directory will contain logs of removed files (if any) and duplicate entries (if using JSON input)
+file_to_url.json: Json file that contains the original URL of each document
+Error Handling
+The script includes error handling for various scenarios, such as:
+Invalid input file/folder paths
+File I/O errors
+Exceptions during data cleaning, splitting, or embedding
+Invalid user input for prompts
+Errors are logged using the logging module
+Notes
+The script assumes that the input JSON data has "url" and "markdown" keys for each entry
+The script uses uuid4 to generate unique IDs for each document in the vector database
+The script processes in batches to deal with a large number of splits
+ + +Disclaimer: This tool is provided "as is" without warranty of any kind. Use it at your own risk. Open source or go away. + +--- + +## ragesink — the last hop into RAGE + +`md-embed.py` ends at **Chroma**, through interactive prompts, defaulting to +`all-MiniLM-L6-v2` (384 dims) or `nomic-embed-text`. RAGE in production runs +**bge-m3 at 1024 dims inside PostgreSQL with pgvectorscale/StreamingDiskANN**. +So the pipeline stopped one hop short of RAGE in three ways at once: wrong store, +wrong dimension, and un-scriptable. + +`ragesink.py` is that hop. No prompts, so it runs from cron or a Makefile: + +```bash +# a whole codebase becomes retrievable +python ../BaseGen/basegen.py ../myrepo -o myrepo.md +python ragesink.py --input myrepo.md --namespace code/myrepo + +python ragesink.py --input ./docs --namespace docs --dry-run # split only +``` + +Dependencies: `psycopg2` and a local Ollama. No langchain, no Chroma. +Full technical reference, with measured costs: [`TECHNICAL.md`](TECHNICAL.md). + +**Re-ingest is cheap.** Chunks carry a content hash (over model ‖ dims ‖ text), +so only what actually changed is re-embedded. Measured on 20 chunks, 4-core CPU: + +| run | embedded | wall | +|---|---|---| +| cold | 20/20 | 515.5 s | +| unchanged | 0/20 | **0.49 s** | +| one section edited | 1/20 | 14.0 s | + +~1050x on an unchanged re-run — which is what makes it safe to put on a timer. +Embedding goes to Ollama `/api/embed` in batches of 32, falling back to the +older per-chunk route on daemons that lack it. + +**Markdown-aware splitting.** Headers first, then packed to 512 words (the +production RAGE chunk size). In a BaseGen document every file is a section, so +the heading carries the path — splitting purely by length would shred a file +across chunks that no longer say where they came from. A `#` inside a fenced +code block is a comment, not a heading, and is treated as such. + +**Index selection is reported, not hidden.** It tries `vectorscale`/DiskANN, +falls back to pgvector/HNSW, then to no index, and prints which it got. A +fallback that hides is just a bug with good manners. + +**Idempotent, and it prunes.** Writes upsert on `(doc_name, chunk_idx)`. An +upsert-only writer never removes anything, so when a document re-chunks SHORTER +the rows above the new count survive holding the previous text — and similarity +search still returns them, served as current. ragesink deletes +`chunk_idx >= len(chunks)` after a successful store, and only after: a run that +stored nothing is no evidence about the document's shape, since the embedder may +simply be down. + +Verified end to end against a live Postgres + Ollama: 20 chunks in, then the +same document truncated → `3/3 stored, pruned 17`. Asking *"how does the +federation handle a peer that does not answer"* returned `federate.py` first, +then `peer.py`. + +### Using it with BaseGen + +BaseGen's stock `config.json` is tuned for code backup, not retrieval, and two +of its defaults work against a RAGE corpus: + +- **`.git` is not excluded** — a repo's git internals land in the corpus as pure + noise (75 KB → 43 KB on a small repo once removed). +- **`*.md` IS excluded** — so READMEs and docs, the highest-value prose in a + repo, are omitted. + +`BaseGen/rage.config.json` is the retrieval-tuned preset: markdown kept, VCS +internals, virtualenvs, `node_modules`, build output and binaries dropped. diff --git a/TECHNICAL.md b/TECHNICAL.md new file mode 100644 index 0000000..37b4e12 --- /dev/null +++ b/TECHNICAL.md @@ -0,0 +1,183 @@ +# TECHNICAL — the RAGE ingestion path + +How a codebase or a folder of documents becomes retrievable by RAGE, what each +stage costs, and which decisions are load-bearing rather than incidental. + +``` +codebase ──BaseGen──▶ markdown ──ragesink──▶ pgvectorscale ──▶ RAGE + │ │ │ │ + │ │ │ └─ StreamingDiskANN, cosine + │ │ └─ split · hash · embed · upsert · prune + │ └─ one document: directory tree + every file + └─ .gitignore + rage.config.json exclusions +``` + +Every number below is measured on the reference box (4-core CPU, no GPU, +PostgreSQL 16, Ollama serving `bge-m3`), not quoted from a datasheet. + +--- + +## 1. BaseGen — codebase to one document + +`basegen.py -o out.md` walks the tree and emits a directory listing +plus every file in a fenced block tagged by language. + +**Its stock config is tuned for code backup, and two defaults actively harm a +retrieval corpus:** + +| default | effect on a RAGE corpus | +|---|---| +| `.git` not excluded | git internals — hooks, objects, sample scripts — become retrievable noise. **75 KB → 43 KB** on a small repo once removed. | +| `*.md` excluded | READMEs and docs are omitted. They are the highest-value prose a repo has. | + +Use `rage.config.json` instead: markdown kept, VCS internals / virtualenvs / +`node_modules` / `__pycache__` / build output / lockfiles / binaries dropped. + +## 2. ragesink — markdown to pgvectorscale + +```bash +python ragesink.py --input out.md --namespace code/myrepo +python ragesink.py --input ./docs --namespace docs --dry-run # split only, no embedding +python ragesink.py --input ./docs --namespace docs --force # re-embed everything +``` + +### 2.1 Splitting — headers first, then 512 words + +Sections are cut on markdown headers, then packed to `--chunk-words` (default +512, the production RAGE chunk size). + +Header-first is not cosmetic. In a BaseGen document **every file is a section**, +so the heading carries the path, and it is prepended to each chunk's text. Split +purely by length and a large file becomes chunks that no longer say where they +came from — the retrieval still "works" and the answers become unattributable. + +A `#` inside a fenced code block is a comment, not a heading. The splitter +tracks fence state, so `# TODO: fix this` in Python does not open a section. + +### 2.2 Identity — the content hash + +``` +sha256(model ‖ dims ‖ chunk_text) +``` + +The model and width are **inside** the hash on purpose. The same text embedded +by a different model is not the same vector, and silently keeping the old row is +how a corpus ends up with two embedding spaces mixed in one table — a failure +that produces plausible-looking nonsense rather than an error. + +### 2.3 Embedding — batched, with a fallback that announces itself + +Chunks go to Ollama `/api/embed` in batches of 32. Older daemons that only speak +`/api/embeddings` fall back to one call per chunk automatically. + +`ensure_schema` tries `vectorscale` → DiskANN, then pgvector → HNSW, then no +index, and **prints which it got**. A fallback that hides is a bug with good +manners; the operator needs to know whether they are running the index they +think they are. + +### 2.4 Writing — upsert, then prune + +Rows upsert on `(doc_name, chunk_idx)`. + +**An upsert-only writer never removes anything.** When a document re-chunks +*shorter*, the rows above the new count survive holding the previous text — and +similarity search still returns them, served as current. So after a store, +`chunk_idx >= len(chunks)` is deleted. + +The guard on that delete is the subtle part. Pruning happens only when the run +has **evidence about the document's shape**: it either wrote something, or it +verified every chunk as already-current. A run that stored nothing because the +embedder was down knows nothing, and must not delete. + +--- + +## 3. Cost, measured + +20 chunks from the RAGEnet codebase, 4-core CPU, no GPU: + +| run | embedded | wall | note | +|---|---|---|---| +| cold ingest | 20 / 20 | **515.5 s** | ~25 s per 512-word chunk; the model dominates | +| re-run, unchanged | 0 / 20 | **0.49 s** | every hash matched | +| re-run, one section edited | 1 / 20 | **14.0 s** | only the changed chunk re-embedded | + +**~1050× on an unchanged re-run.** That is the number that decides whether this +can go on a timer, and it comes from the content hash, not from batching. + +Batching cuts request round-trips, but on CPU the model's forward pass is +essentially all of the cost, so it barely moves the cold number. Say so plainly +rather than claiming a speedup that a GPU would show and this box does not. + +The practical consequence: **the first ingest of a large corpus is an overnight +job; every ingest after it is seconds.** Budget accordingly — a 1,000-chunk +corpus is roughly 7 hours cold on this hardware. + +--- + +## 4. Schema + +```sql +CREATE EXTENSION IF NOT EXISTS vector; +CREATE TABLE doc_embeddings ( + doc_name TEXT NOT NULL, + chunk_idx INTEGER NOT NULL, + text_content TEXT NOT NULL, + embedding vector(1024), -- bge-m3 + content_hash TEXT, -- added idempotently; enables the skip + PRIMARY KEY (doc_name, chunk_idx) +); +CREATE INDEX doc_embeddings_diskann ON doc_embeddings + USING diskann (embedding vector_cosine_ops); +``` + +`ALTER TABLE ... ADD COLUMN IF NOT EXISTS content_hash` runs on every start, so +ragesink attaches to an existing RAGE store with no migration. Chunks written +before the column existed have a NULL hash and are re-embedded once. + +**`doc_name` is the namespace.** `--namespace code/myrepo` yields +`code/myrepo/src/thing.py`, which is what retrieval filters and excludes on. +Gated subtrees must be excluded **in SQL**, not filtered after the fact: filter +afterwards and private chunks still consume slots inside the top-k, silently +displacing public results. Nothing leaks, the answer just quietly gets worse and +nothing logs it. + +--- + +## 5. Retrieval — what this buys + +```sql +SELECT doc_name, chunk_idx, 1 - (embedding <=> $1::vector) AS similarity +FROM doc_embeddings +WHERE doc_name LIKE 'code/myrepo/%' +ORDER BY embedding <=> $1::vector +LIMIT 8; +``` + +Verified against the ingested RAGEnet corpus — *"how does the federation handle +a peer that does not answer"*: + +``` +sim=0.6119 federate.py ← the file that implements exactly that +sim=0.5703 peer.py +sim=0.5660 examples/federate_demo.py +``` + +No keyword in that query appears in those filenames. That is the whole point. + +## 6. Federating it + +One store is a node. [RAGEnet](https://github.com/GATERAGE/RAGEnet) queries many +and fuses by rank (RRF), so a pgvectorscale store, a WordPress index and a JSONL +file answer together without any corpus moving. + +## 7. Operational notes + +- **Re-ingest is safe and cheap.** Put it on a timer; unchanged corpora cost + under a second. +- **`--dry-run` first** on a new corpus — it reports the chunk count, and chunk + count × ~25 s is the cold budget on CPU. +- **Never mix embedders in one table.** The width is in the schema and the model + is in the hash, so changing either re-embeds rather than corrupting silently. +- **A failed embed is skipped, not faked.** The chunk is left absent and the + count reported, so a partial ingest is visible instead of a corpus with holes + that nothing records. diff --git a/md-embed.py b/md-embed.py index 7bc0114..50ca57b 100644 --- a/md-embed.py +++ b/md-embed.py @@ -1,3 +1,5 @@ +# mdmbed (c) 2005 w3d + import json import re import os diff --git a/ragesink.py b/ragesink.py new file mode 100644 index 0000000..89b9572 --- /dev/null +++ b/ragesink.py @@ -0,0 +1,426 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""ragesink — land markdown in RAGE's actual store: pgvectorscale. + +WHY THIS EXISTS + md-embed.py ends at Chroma, through interactive input() prompts, defaulting + to all-MiniLM-L6-v2 (384 dims) or nomic-embed-text. RAGE in production runs + **bge-m3 at 1024 dims inside PostgreSQL with pgvectorscale/StreamingDiskANN**. + So the pipeline + + codebase --BaseGen--> markdown --md-embed--> chunks --> Chroma + + stops one hop short of RAGE, in three ways at once: wrong store, wrong + dimension, and un-scriptable. ragesink is that last hop, and it takes no + prompts so it can run from cron or a Makefile. + + python basegen.py --codebase ../myrepo --output myrepo.md + python ragesink.py --input myrepo.md --namespace code/myrepo + + Dependencies: psycopg2 and a local Ollama. No langchain, no Chroma. + +IDEMPOTENCE — the part that is easy to get wrong + Writes upsert on (doc_name, chunk_idx). An upsert-only writer NEVER REMOVES + anything, so when a document re-chunks SHORTER the rows above the new count + survive holding the previous text — and they are still returned by + similarity search, served as current. ragesink prunes chunk_idx >= the new + count after a successful store, and only after: a run that stored nothing is + no evidence about the document's shape (the embedder may simply be down). +""" +from __future__ import annotations + +import argparse +import hashlib +import time +import json +import os +import re +import sys +import urllib.request +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Tuple + +DEFAULT_DSN = os.environ.get("MINDX_DB_DSN", "postgresql://mindx:mindx@localhost:5432/mindx") +DEFAULT_OLLAMA = os.environ.get("OLLAMA_URL", "http://localhost:11434") +DEFAULT_MODEL = os.environ.get("RAGE_EMBED_MODEL", "bge-m3") +DEFAULT_DIMS = int(os.environ.get("RAGE_EMBED_DIMS", "1024")) +CHUNK_WORDS = 512 # matches the production RAGE chunker +MIN_CHUNK_CHARS = 50 + +DDL = """ +CREATE EXTENSION IF NOT EXISTS vector; +CREATE TABLE IF NOT EXISTS doc_embeddings ( + doc_name TEXT NOT NULL, + chunk_idx INTEGER NOT NULL, + text_content TEXT NOT NULL, + embedding vector({dims}), + PRIMARY KEY (doc_name, chunk_idx) +); +""" +# Additive and idempotent, so this runs against an existing RAGE store without a +# migration. The hash is what makes a re-ingest cheap: unchanged chunks are never +# re-embedded, and embedding is ~100% of the cost. +DDL_HASH = "ALTER TABLE doc_embeddings ADD COLUMN IF NOT EXISTS content_hash TEXT;" + +BATCH = 32 # chunks per /api/embed call +RECONNECT_TRIES = 4 # a long tunnelled ingest will see the tunnel flap + + +class SinkGone(RuntimeError): + """The store went away mid-run.""" + + +def alive(cur) -> bool: + """Is the connection still usable? + + Checked BEFORE each embed batch, not after. An embed costs ~25 s of CPU; a + dropped tunnel that is only noticed at INSERT time means every one of those + seconds is spent producing a vector with nowhere to go. Learned the hard + way, twice, on long ingests over an SSH tunnel.""" + try: + cur.execute("SELECT 1") + cur.fetchone() + return True + except Exception: + return False + + +def chunk_hash(text: str, model: str, dims: int) -> str: + """Identity of an embedded chunk. Includes the model and width: the same text + under a different embedder is NOT the same vector, and silently keeping the + old one is how a corpus ends up with mixed embedding spaces.""" + h = hashlib.sha256() + h.update(model.encode()); h.update(b"\x00"); h.update(str(dims).encode()) + h.update(b"\x00"); h.update(text.encode("utf-8", "replace")) + return h.hexdigest() + +# ── markdown-aware splitting ───────────────────────────────────────────────── +HEADER = re.compile(r"^(#{1,6})\s+(.*)$", re.M) +FENCE = re.compile(r"^```") + + +def split_markdown(text: str, chunk_words: int = CHUNK_WORDS) -> List[Tuple[str, str]]: + """Split on headers first, then pack to chunk_words. Returns (heading, text). + + Header-first matters for a BaseGen document: every file is a section, so the + heading carries the path. Splitting purely by length would shred a file + across chunks that no longer say which file they came from. Fenced code + blocks are never split on a '#' inside them — a comment is not a heading. + """ + lines = text.splitlines() + sections: List[Tuple[str, List[str]]] = [] + cur_head: str = "" + cur: List[str] = [] + in_fence = False + for ln in lines: + if FENCE.match(ln.strip()): + in_fence = not in_fence + m = None if in_fence else HEADER.match(ln) + if m: + if cur: + sections.append((cur_head, cur)) + cur_head, cur = m.group(2).strip(), [] + else: + cur.append(ln) + if cur or cur_head: + sections.append((cur_head, cur)) + + # Pack consecutive sections up to chunk_words instead of emitting one chunk + # per section. A doc with many short sections would otherwise produce a + # chunk per heading — on the mindX docs that was 440 chunks averaging 174 + # words against a 512 cap: 3x the embedding cost, and a corpus fragmented + # into pieces too small to carry their own context. Oversized sections still + # split, and a split section repeats its heading on every piece. + out: List[Tuple[str, str]] = [] + buf_head: str = "" + buf: List[str] = [] + buf_n = 0 + + def flush() -> None: + nonlocal buf, buf_n, buf_head + if buf: + piece = " ".join(buf) + if len(piece.strip()) > MIN_CHUNK_CHARS: + out.append((buf_head, piece)) + buf, buf_n = [], 0 + + for head, body in sections: + words = "\n".join(body).split() + if not words: + continue + if len(words) >= chunk_words: # big section: flush, then split it + flush() + for i in range(0, len(words), chunk_words): + piece = " ".join(words[i:i + chunk_words]) + if len(piece.strip()) > MIN_CHUNK_CHARS: + out.append((head, piece)) + buf_head = "" + continue + if buf_n + len(words) > chunk_words: # would overflow: close this chunk + flush() + if not buf: + buf_head = head + buf.extend(words) + buf_n += len(words) + flush() + return out + + +# ── embedding ──────────────────────────────────────────────────────────────── +def _post(url: str, payload: dict, timeout: float): + req = urllib.request.Request(url, data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read().decode()) + + +def embed(text: str, url: str = DEFAULT_OLLAMA, model: str = DEFAULT_MODEL, + timeout: float = 120.0) -> Optional[List[float]]: + """One chunk. Kept for callers and for the per-chunk fallback path.""" + out = embed_batch([text], url=url, model=model, timeout=timeout) + return out[0] if out and out[0] else None + + +def embed_batch(texts: List[str], url: str = DEFAULT_OLLAMA, model: str = DEFAULT_MODEL, + timeout: float = 600.0) -> List[Optional[List[float]]]: + """Embed many chunks in ONE call via /api/embed. + + This is the difference between a usable ingest and an unusable one: a + per-chunk POST pays the request round trip and the model's per-call overhead + on every chunk. Falls back to the older per-chunk /api/embeddings when the + daemon does not speak /api/embed, so an older Ollama still works — slower, + and it says so. + """ + if not texts: + return [] + clipped = [t[:8000] for t in texts] + try: + d = _post(url.rstrip("/") + "/api/embed", {"model": model, "input": clipped}, timeout) + vecs = d.get("embeddings") + if isinstance(vecs, list) and len(vecs) == len(clipped): + return [v if v else None for v in vecs] + except Exception: + pass # fall through to the old route + out: List[Optional[List[float]]] = [] + for t in clipped: + try: + d = _post(url.rstrip("/") + "/api/embeddings", {"model": model, "prompt": t}, timeout) + out.append(d.get("embedding") or None) + except Exception as exc: # noqa: BLE001 + print(f" ! embed failed: {type(exc).__name__}: {exc}", file=sys.stderr) + out.append(None) + return out + + +# ── the sink ───────────────────────────────────────────────────────────────── +def has_hash_column(cur) -> bool: + cur.execute("SELECT 1 FROM information_schema.columns " + "WHERE table_name='doc_embeddings' AND column_name='content_hash'") + return cur.fetchone() is not None + + +def ensure_schema(cur, dims: int) -> str: + cur.execute(DDL.format(dims=dims)) + # Attaching to a store you do not OWN is normal — a shared RAGE database + # hands out INSERT/UPDATE without table ownership, and ALTER then fails with + # InsufficientPrivilege. Degrade instead of dying: without content_hash the + # skip optimisation is off and every run re-embeds, which is slow but + # correct. Announce it, because a silent 100x slowdown looks like a hang. + try: + cur.execute(DDL_HASH) + except Exception: + cur.connection.rollback() + cur.execute(DDL.format(dims=dims)) + try: + cur.execute("CREATE EXTENSION IF NOT EXISTS vectorscale;") + cur.execute("CREATE INDEX IF NOT EXISTS doc_embeddings_diskann " + "ON doc_embeddings USING diskann (embedding vector_cosine_ops);") + return "pgvectorscale/diskann" + except Exception: + # Graceful degradation, reported rather than hidden: a fallback that + # hides is just a bug with good manners. + cur.connection.rollback() + cur.execute(DDL.format(dims=dims)) + try: + cur.execute("CREATE INDEX IF NOT EXISTS doc_embeddings_hnsw " + "ON doc_embeddings USING hnsw (embedding vector_cosine_ops);") + return "pgvector/hnsw" + except Exception: + cur.connection.rollback() + cur.execute(DDL.format(dims=dims)) + return "pgvector/no-index" + + +def sink(doc_name: str, chunks: List[Tuple[str, str]], cur, *, url: str, model: str, + dims: int, force: bool = False, hashing: bool = True) -> Tuple[int, int, int]: + """Embed and upsert. Returns (stored, skipped, pruned). + + Only chunks whose content hash changed are embedded. Embedding is ~100% of + the cost of an ingest, so a re-run over an unchanged corpus does almost no + work — which is what makes this safe to put on a timer. + """ + bodies = [(f"{h}\n\n{t}" if h else t) for h, t in chunks] + hashes = [chunk_hash(b, model, dims) for b in bodies] + + have: Dict[int, str] = {} + if not force and hashing: + cur.execute("SELECT chunk_idx, content_hash FROM doc_embeddings WHERE doc_name=%s", (doc_name,)) + have = {i: hsh for i, hsh in cur.fetchall() if hsh} + + todo = [i for i, hsh in enumerate(hashes) if have.get(i) != hsh] + skipped = len(chunks) - len(todo) + + stored = 0 + for start in range(0, len(todo), BATCH): + idxs = todo[start:start + BATCH] + if not alive(cur): + raise SinkGone(f"store unreachable before embedding {len(idxs)} chunk(s) " + f"of {doc_name}; stopping rather than embedding into the void") + vecs = embed_batch([bodies[i] for i in idxs], url=url, model=model) + for i, vec in zip(idxs, vecs): + if vec is None: + continue + if hashing: + cur.execute( + "INSERT INTO doc_embeddings (doc_name, chunk_idx, text_content, embedding, content_hash) " + "VALUES (%s,%s,%s,%s::vector,%s) " + "ON CONFLICT (doc_name, chunk_idx) DO UPDATE " + "SET text_content=EXCLUDED.text_content, embedding=EXCLUDED.embedding, " + " content_hash=EXCLUDED.content_hash", + (doc_name, i, bodies[i], str(vec), hashes[i]), + ) + else: + cur.execute( + "INSERT INTO doc_embeddings (doc_name, chunk_idx, text_content, embedding) " + "VALUES (%s,%s,%s,%s::vector) " + "ON CONFLICT (doc_name, chunk_idx) DO UPDATE " + "SET text_content=EXCLUDED.text_content, embedding=EXCLUDED.embedding", + (doc_name, i, bodies[i], str(vec)), + ) + stored += 1 + + pruned = 0 + # Prune only when this run has evidence about the document's shape: it either + # wrote something, or it verified every chunk as already-current. A run that + # stored nothing because the embedder was down must not delete anything. + if stored or (skipped == len(chunks) and chunks): + cur.execute("DELETE FROM doc_embeddings WHERE doc_name=%s AND chunk_idx >= %s", + (doc_name, len(chunks))) + pruned = cur.rowcount or 0 + return stored, skipped, pruned + + +def iter_inputs(path: Path) -> Iterable[Path]: + if path.is_dir(): + yield from sorted(p for p in path.rglob("*.md") if p.is_file()) + else: + yield path + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description="land markdown in pgvectorscale for RAGE") + ap.add_argument("--input", required=True, type=Path, help="markdown file or folder") + ap.add_argument("--namespace", default="", help="doc_name prefix, e.g. code/myrepo") + ap.add_argument("--dsn", default=DEFAULT_DSN) + ap.add_argument("--ollama", default=DEFAULT_OLLAMA) + ap.add_argument("--model", default=DEFAULT_MODEL) + ap.add_argument("--dims", type=int, default=DEFAULT_DIMS) + ap.add_argument("--chunk-words", type=int, default=CHUNK_WORDS) + ap.add_argument("--strip-ext", action="store_true", + help="drop the .md from doc_name, to match an existing store's convention") + ap.add_argument("--force", action="store_true", help="re-embed even unchanged chunks") + ap.add_argument("--dry-run", action="store_true", help="split and report; embed nothing") + a = ap.parse_args(argv) + + files = list(iter_inputs(a.input)) + if not files: + print("no markdown found", file=sys.stderr); return 1 + + if a.dry_run: + total = 0 + for f in files: + ch = split_markdown(f.read_text(encoding="utf-8", errors="replace"), a.chunk_words) + total += len(ch) + print(f" {len(ch):4d} chunks {f}") + print(f"dry run: {len(files)} file(s), {total} chunks, model={a.model} dims={a.dims}") + return 0 + + global psycopg2 + try: + import psycopg2 + except ImportError: + print("psycopg2 required: pip install psycopg2-binary", file=sys.stderr); return 2 + + conn = psycopg2.connect(a.dsn); conn.autocommit = False + cur = conn.cursor() + backend = ensure_schema(cur, a.dims); conn.commit() + hashing = has_hash_column(cur) + print(f"[ragesink] backend: {backend} · model {a.model} @ {a.dims} dims" + + ("" if hashing else " · NO content_hash column (not owner): re-embedding everything"), + flush=True) + + files_done = chunks_done = pruned_total = skipped_total = 0 + for f in files: + rel = f.name if a.input.is_file() else str(f.relative_to(a.input)) + if a.strip_ext and rel.lower().endswith(".md"): + rel = rel[:-3] # attach to a store that names docs without the extension + doc_name = f"{a.namespace.rstrip('/')}/{rel}" if a.namespace else rel + chunks = split_markdown(f.read_text(encoding="utf-8", errors="replace"), a.chunk_words) + if not chunks: + print(f" --- {doc_name}: no chunks", flush=True); continue + # A long ingest over an SSH tunnel WILL see the tunnel flap. The + # pre-check catches a connection already dead; this catches one that + # dies mid-write. Reconnect and retry the file — files commit + # atomically, so a retry is clean, and already-committed files are + # skipped by hash on the way back through. + stored = skipped = pruned = 0 + for attempt in range(RECONNECT_TRIES + 1): + try: + stored, skipped, pruned = sink(doc_name, chunks, cur, url=a.ollama, model=a.model, + dims=a.dims, force=a.force, hashing=hashing) + break + except (psycopg2.OperationalError, SinkGone) as exc: + if attempt >= RECONNECT_TRIES: + print(f"\n[ragesink] ABORTED after {RECONNECT_TRIES} reconnect attempts: {exc}", + file=sys.stderr, flush=True) + print(f"[ragesink] {files_done} file(s) committed. Re-run to resume — " + f"stored chunks are skipped by hash.", file=sys.stderr, flush=True) + try: conn.close() + except Exception: pass + return 3 + wait = 5 * (attempt + 1) + print(f" ~ store dropped ({type(exc).__name__}); reconnecting in {wait}s " + f"[{attempt+1}/{RECONNECT_TRIES}]", file=sys.stderr, flush=True) + time.sleep(wait) + try: conn.close() + except Exception: pass + try: + conn = psycopg2.connect(a.dsn); conn.autocommit = False; cur = conn.cursor() + except Exception as e2: + print(f" ~ reconnect failed: {e2}", file=sys.stderr, flush=True) + try: + conn.commit() + except psycopg2.OperationalError: + print(f" ~ commit lost for {doc_name}; it will be redone on the next run", + file=sys.stderr, flush=True) + try: conn.close() + except Exception: pass + conn = psycopg2.connect(a.dsn); conn.autocommit = False; cur = conn.cursor() + continue + files_done += 1; chunks_done += stored; pruned_total += pruned; skipped_total += skipped + failed = len(chunks) - stored - skipped + bits = [] + if skipped: bits.append(f"{skipped} unchanged") + if pruned: bits.append(f"pruned {pruned}") + if failed: bits.append(f"{failed} embed failures") + print(f" {stored:4d}/{len(chunks):<4d} {doc_name}" + (" (" + ", ".join(bits) + ")" if bits else ""), + flush=True) + + cur.close(); conn.close() + print(f"[ragesink] {files_done} file(s) · {chunks_done} embedded · " + f"{skipped_total} unchanged · {pruned_total} orphan(s) pruned") + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main())