diff --git a/README.md b/README.md index 26d5241..2965e9f 100644 --- a/README.md +++ b/README.md @@ -135,17 +135,28 @@ buffers are reconstructed rather than serialized. The service wraps that binary snapshot with deterministic JSON metadata in a versioned `RAGSTATE` container and a BLAKE2b checksum. -Each persistent insertion is staged against a cloned index while holding the store -lock. The complete container is written to a sibling temporary file, flushed and +Each persistent corpus mutation is staged while holding the store lock. The +complete container is written to a sibling temporary file, flushed and `fsync`ed, atomically replaced, and its parent directory is `fsync`ed before the new in-memory state is published and HTTP success is returned. A failed commit -therefore retains the previous file and previous live corpus. Corrupt state causes +therefore retains the previous file and previous live corpus. This transactional +stage → persist → publish rule applies equally to add, delete, and replace. +Corrupt state causes startup to fail; it is never silently discarded. Stored index settings, embedding dimension, backend class, and model name (where applicable) must match the runtime. +Deletion and replacement compact the corpus with a full deterministic HNSW rebuild, +ordered by ascending document id and then chunk ordinal. Surviving chunks reuse +their vectors through the index `vector(id)` API; only replacement text is embedded. +Document ids are stable (and deleted ids are never reused), while internal chunk ids +may be reassigned so they remain dense and exactly aligned with HNSW vector ids. +This rebuild is **O(number of live vectors)** and is deliberately chosen for +correctness and simplicity at the current portfolio/demo scale, not as +production-scale constant-time deletion. + Limitations: snapshots coordinate one service process only; do not point multiple workers or hosts at the same file. Durability ultimately depends on the filesystem's -`fsync` and atomic-replace semantics. There is no deletion, migration between +`fsync` and atomic-replace semantics. There is no migration between embedder configurations, or automatic recovery of a corrupt snapshot. ### Public API hardening @@ -153,7 +164,7 @@ embedder configurations, or automatic recovery of a corrupt snapshot. | Variable | Default | Behavior | |----------|---------|----------| | `RAG_STATE_PATH` | empty | Enables the versioned durable corpus snapshot at the configured file. Empty preserves in-memory behavior. | -| `RAG_UPLOADS_ENABLED` | `0` | Enables `POST /documents` only when explicitly set to `1`, `true`, `yes`, or `on`. | +| `RAG_UPLOADS_ENABLED` | `0` | Enables corpus mutations (`POST`, `PUT`, and `DELETE /documents`) only when explicitly set to `1`, `true`, `yes`, or `on`. | | `RAG_MAX_TITLE_CHARS` | `200` | Maximum document title length. | | `RAG_MAX_DOCUMENT_CHARS` | `1000000` | Maximum document text length. | | `RAG_MAX_QUESTION_CHARS` | `2000` | Maximum query length. | @@ -165,7 +176,7 @@ embedder configurations, or automatic recovery of a corrupt snapshot. Chunking accepts `max_words` from 1 through 1000 and `overlap` from 0 through 999, with overlap strictly smaller than the chunk size. Query `k` remains 1 through 50 and `ef_search` is limited to 1 through 2000. Oversized -bodies return `413`, invalid fields return `422`, disabled uploads return +bodies return `413`, invalid fields return `422`, disabled corpus mutations return `403`, and Claude timeouts return `504`. The Next.js same-origin proxy preserves these status codes and service messages. diff --git a/app/app/api/documents/[id]/route.ts b/app/app/api/documents/[id]/route.ts new file mode 100644 index 0000000..c070cf5 --- /dev/null +++ b/app/app/api/documents/[id]/route.ts @@ -0,0 +1,44 @@ +import { NextResponse } from "next/server"; +import { deleteDocument, replaceDocument } from "@/app/lib/rag"; +import { INPUT_LIMITS, readLimitedJson } from "@/app/lib/limits"; +import { apiErrorResponse } from "@/app/lib/api-errors"; + +type Context = { params: Promise<{ id: string }> }; + +function documentId(raw: string): number | null { + const id = Number(raw); + return Number.isSafeInteger(id) && id >= 0 ? id : null; +} + +export async function DELETE(_: Request, context: Context) { + const id = documentId((await context.params).id); + if (id === null) return NextResponse.json({ error: "invalid document id" }, { status: 422 }); + try { + return NextResponse.json(await deleteDocument(id)); + } catch (err) { + return apiErrorResponse(err); + } +} + +export async function PUT(req: Request, context: Context) { + const id = documentId((await context.params).id); + if (id === null) return NextResponse.json({ error: "invalid document id" }, { status: 422 }); + try { + const body = await readLimitedJson<{ title?: unknown; text?: unknown }>(req); + if (typeof body.title !== "string" || !body.title.trim()) { + return NextResponse.json({ error: "title is required" }, { status: 422 }); + } + if (body.title.trim().length > INPUT_LIMITS.titleChars) { + return NextResponse.json({ error: `title must be at most ${INPUT_LIMITS.titleChars} characters` }, { status: 422 }); + } + if (typeof body.text !== "string" || !body.text.trim()) { + return NextResponse.json({ error: "document text is required" }, { status: 422 }); + } + if (body.text.length > INPUT_LIMITS.documentChars) { + return NextResponse.json({ error: `document text must be at most ${INPUT_LIMITS.documentChars} characters` }, { status: 422 }); + } + return NextResponse.json(await replaceDocument(id, body.title.trim(), body.text)); + } catch (err) { + return apiErrorResponse(err); + } +} diff --git a/app/app/globals.css b/app/app/globals.css index 732c4cd..862ec13 100644 --- a/app/app/globals.css +++ b/app/app/globals.css @@ -165,6 +165,21 @@ button { font-weight: 700; } +.text-button:disabled { + cursor: not-allowed; + opacity: 0.45; +} + +.danger-button { + color: var(--danger); +} + +.document-row > .document-actions { + display: flex; + gap: 12px; + margin-left: auto; +} + .metrics-grid { display: grid; grid-template-columns: repeat(4, 1fr); diff --git a/app/app/lib/rag.ts b/app/app/lib/rag.ts index a4f48bb..5af0c02 100644 --- a/app/app/lib/rag.ts +++ b/app/app/lib/rag.ts @@ -124,6 +124,21 @@ export function addDocument(title: string, text: string): Promise }); } +export function deleteDocument(id: number): Promise { + return call(`/documents/${id}`, { method: "DELETE" }); +} + +export function replaceDocument( + id: number, + title: string, + text: string, +): Promise { + return call(`/documents/${id}`, { + method: "PUT", + body: JSON.stringify({ title, text }), + }); +} + export function listDocuments(): Promise { return call("/documents"); } diff --git a/app/app/page.tsx b/app/app/page.tsx index 354fa61..05e6873 100644 --- a/app/app/page.tsx +++ b/app/app/page.tsx @@ -68,6 +68,8 @@ export default function Home() { const [uploadError, setUploadError] = useState(null); const [uploadNotice, setUploadNotice] = useState(null); const [fileInputKey, setFileInputKey] = useState(0); + const [editingDocumentId, setEditingDocumentId] = useState(null); + const [mutationDocumentId, setMutationDocumentId] = useState(null); const uploadsEnabled = stats?.uploads_enabled ?? false; const titleLimit = stats?.limits.title_chars ?? 200; @@ -149,20 +151,26 @@ export default function Home() { setUploadError(null); setUploadNotice(null); try { - const created = await requestJson("/api/documents", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - title: documentTitle.trim(), - text: documentText, - }), - }); + const created = await requestJson( + editingDocumentId === null + ? "/api/documents" + : `/api/documents/${editingDocumentId}`, + { + method: editingDocumentId === null ? "POST" : "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + title: documentTitle.trim(), + text: documentText, + }), + }, + ); setUploadNotice( - `Indexed “${created.title}” as ${created.n_chunks} ${created.n_chunks === 1 ? "chunk" : "chunks"}.`, + `${editingDocumentId === null ? "Indexed" : "Replaced"} “${created.title}” as ${created.n_chunks} ${created.n_chunks === 1 ? "chunk" : "chunks"}.`, ); setDocumentTitle(""); setDocumentText(""); setFileInputKey((key) => key + 1); + setEditingDocumentId(null); await loadWorkspace(); } catch (err) { setUploadError(err instanceof Error ? err.message : "Unable to index document"); @@ -171,6 +179,39 @@ export default function Home() { } } + function beginReplacement(document: DocumentInfo) { + setEditingDocumentId(document.id); + setDocumentTitle(document.title); + setDocumentText(""); + setUploadError(null); + setUploadNotice(`Paste the replacement text for “${document.title}”.`); + } + + async function removeDocument(document: DocumentInfo) { + if ( + !uploadsEnabled || + !window.confirm(`Delete “${document.title}” and all of its chunks?`) + ) + return; + setMutationDocumentId(document.id); + setWorkspaceError(null); + try { + await requestJson(`/api/documents/${document.id}`, { + method: "DELETE", + }); + if (editingDocumentId === document.id) { + setEditingDocumentId(null); + setDocumentTitle(""); + setDocumentText(""); + } + await loadWorkspace(); + } catch (err) { + setWorkspaceError(err instanceof Error ? err.message : "Unable to delete document"); + } finally { + setMutationDocumentId(null); + } + } + const hashedRetrieval = stats?.embedder === "HashedEmbedder"; const retrievalLabel = stats ? hashedRetrieval @@ -356,8 +397,14 @@ export default function Home() { >
-

ADD SOURCE

-

Index a document

+

+ {editingDocumentId === null ? "ADD SOURCE" : "REPLACE SOURCE"} +

+

+ {editingDocumentId === null + ? "Index a document" + : `Replace document ${editingDocumentId}`} +

{stats ? (uploadsEnabled ? "Enabled" : "Locked") : "Loading"} @@ -429,7 +476,11 @@ export default function Home() { !documentText.trim() } > - {uploadLoading ? "Indexing…" : "Add to index"} + {uploadLoading + ? "Saving…" + : editingDocumentId === null + ? "Add to index" + : "Replace document"}
@@ -470,7 +521,24 @@ export default function Home() { {document.n_chunks} {document.n_chunks === 1 ? "chunk" : "chunks"} - Indexed +
+ + +
))} diff --git a/service/rag_service/app.py b/service/rag_service/app.py index 12b4464..fc8595e 100644 --- a/service/rag_service/app.py +++ b/service/rag_service/app.py @@ -4,6 +4,7 @@ talks to this. Endpoints: POST /documents upload text → chunk → embed → insert into HNSW + PUT/DELETE /documents/{id} replace or delete through HNSW compaction POST /query embed question → HNSW search → Claude answer + sources GET /stats index size / config GET /documents list indexed documents @@ -28,7 +29,7 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send from .generation import GenerationTimeoutError, generate_answer -from .store import DocumentStore +from .store import DocumentNotFoundError, DocumentStore from hnsw_rag import get_embedder log = logging.getLogger(__name__) @@ -79,7 +80,7 @@ def _add_cors_middleware(target: FastAPI, origins: List[str]) -> None: target.add_middleware( CORSMiddleware, allow_origins=origins, - allow_methods=["GET", "POST", "OPTIONS"], + allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"], allow_headers=["Content-Type"], ) @@ -305,6 +306,38 @@ def add_document(doc: DocumentIn) -> DocumentOut: return DocumentOut(id=result.id, title=result.title, n_chunks=result.n_chunks) +@app.delete( + "/documents/{document_id}", + response_model=DocumentOut, + dependencies=[Depends(_require_uploads_enabled)], +) +def delete_document(document_id: int) -> DocumentOut: + try: + result = _store.delete_document(document_id) + except DocumentNotFoundError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + return DocumentOut(id=result.id, title=result.title, n_chunks=result.n_chunks) + + +@app.put( + "/documents/{document_id}", + response_model=DocumentOut, + dependencies=[Depends(_require_uploads_enabled)], +) +def replace_document(document_id: int, doc: DocumentIn) -> DocumentOut: + try: + result = _store.replace_document( + document_id, + doc.title, + doc.text, + max_words=doc.max_words, + overlap=doc.overlap, + ) + except DocumentNotFoundError as exc: + raise HTTPException(status_code=404, detail=str(exc)) from exc + return DocumentOut(id=result.id, title=result.title, n_chunks=result.n_chunks) + + @app.post("/query", response_model=QueryOut) def query(q: QueryIn) -> QueryOut: chunks = _store.retrieve(q.question, k=q.k, ef_search=q.ef_search) diff --git a/service/rag_service/store.py b/service/rag_service/store.py index ff560f0..7cb99a5 100644 --- a/service/rag_service/store.py +++ b/service/rag_service/store.py @@ -50,6 +50,10 @@ class RetrievedChunk: score: float +class DocumentNotFoundError(LookupError): + """Raised when a corpus mutation targets an unknown document id.""" + + def _embedder_identity(embedder: Embedder) -> dict: identity = {"type": type(embedder).__name__, "dim": embedder.dim} model_name = getattr(embedder, "model_name", None) @@ -282,6 +286,107 @@ def add_document( self._next_doc_id = next_doc_id return doc + def _rebuild( + self, + documents: Dict[int, Document], + replacement: Optional[tuple[int, List[Chunk], List[List[float]]]] = None, + ) -> tuple[Hnsw, Dict[int, StoredChunk], Dict[int, Document]]: + """Build a dense deterministic graph while reusing every survivor vector. + + The caller holds ``_lock``. Chunks are ordered by document id and ordinal, + making the rebuilt HNSW id equal to the new dense chunk metadata id. + """ + replacement_doc_id = replacement[0] if replacement is not None else None + replacement_chunks = replacement[1] if replacement is not None else [] + replacement_vectors = replacement[2] if replacement is not None else [] + live: List[tuple[int, str, int, List[float]]] = [] + for doc_id in sorted(documents): + document = documents[doc_id] + if doc_id == replacement_doc_id: + live.extend( + (doc_id, chunk.text, chunk.index, vector) + for chunk, vector in zip(replacement_chunks, replacement_vectors) + ) + continue + old_chunks = sorted( + (chunk for chunk in self._chunks.values() if chunk.doc_id == doc_id), + key=lambda chunk: chunk.ordinal, + ) + live.extend( + (doc_id, chunk.text, chunk.ordinal, self._index.vector(chunk.id)) + for chunk in old_chunks + ) + + index = self._new_index() + ids = index.insert_batch([item[3] for item in live]) if live else [] + if ids != list(range(len(live))): + raise RuntimeError("rebuilt index did not assign dense chunk ids") + chunks = { + chunk_id: StoredChunk( + chunk_id, text, doc_id, documents[doc_id].title, ordinal + ) + for chunk_id, (doc_id, text, ordinal, _) in enumerate(live) + } + counts = {doc_id: 0 for doc_id in documents} + for chunk in chunks.values(): + counts[chunk.doc_id] += 1 + rebuilt_documents = { + doc_id: Document(doc_id, document.title, counts[doc_id]) + for doc_id, document in documents.items() + } + return index, chunks, rebuilt_documents + + def delete_document(self, doc_id: int) -> Document: + with self._lock: + try: + deleted = self._documents[doc_id] + except KeyError as exc: + raise DocumentNotFoundError(f"document {doc_id} was not found") from exc + documents = dict(self._documents) + del documents[doc_id] + index, chunks, documents = self._rebuild(documents) + if self.state_path is not None: + self._persist(index, chunks, documents, self._next_doc_id) + self._index, self._chunks, self._documents = index, chunks, documents + return deleted + + def replace_document( + self, + doc_id: int, + title: str, + text: str, + *, + max_words: int = 180, + overlap: int = 40, + ) -> Document: + with self._lock: + if doc_id not in self._documents: + raise DocumentNotFoundError(f"document {doc_id} was not found") + replacement_chunks = chunk_text( + text, max_words=max_words, overlap=overlap, source=title + ) + vectors = ( + self.embedder.embed([chunk.text for chunk in replacement_chunks]) + if replacement_chunks + else [] + ) + if len(vectors) != len(replacement_chunks): + raise ValueError( + f"embedder returned {len(vectors)} vectors for {len(replacement_chunks)} chunks" + ) + with self._lock: + if doc_id not in self._documents: + raise DocumentNotFoundError(f"document {doc_id} was not found") + documents = dict(self._documents) + documents[doc_id] = Document(doc_id, title, len(replacement_chunks)) + index, chunks, documents = self._rebuild( + documents, (doc_id, replacement_chunks, vectors) + ) + if self.state_path is not None: + self._persist(index, chunks, documents, self._next_doc_id) + self._index, self._chunks, self._documents = index, chunks, documents + return documents[doc_id] + def retrieve( self, query: str, k: int = 5, ef_search: int = 100 ) -> List[RetrievedChunk]: diff --git a/service/tests/test_document_lifecycle.py b/service/tests/test_document_lifecycle.py new file mode 100644 index 0000000..d43a2a3 --- /dev/null +++ b/service/tests/test_document_lifecycle.py @@ -0,0 +1,132 @@ +import pytest + +from hnsw_rag.embeddings import HashedEmbedder +from rag_service.store import DocumentNotFoundError, DocumentStore + + +class CountingEmbedder: + def __init__(self, dim=32): + self.inner = HashedEmbedder(dim=dim) + self.dim = dim + self.calls = [] + + def embed(self, texts): + self.calls.append(list(texts)) + return self.inner.embed(texts) + + +def assert_dense(store): + assert len(store._index) == len(store._chunks) + assert set(store._chunks) == set(range(len(store._index))) + assert all(chunk.doc_id in store._documents for chunk in store._chunks.values()) + for document in store.documents(): + assert document.n_chunks == sum( + chunk.doc_id == document.id for chunk in store._chunks.values() + ) + + +def test_delete_rebuilds_dense_ids_without_reembedding_and_keeps_doc_ids_monotonic(): + embedder = CountingEmbedder() + store = DocumentStore(embedder, min_score=None) + alpha = store.add_document("A", "alpha-specific content") + beta = store.add_document("B", "beta-specific content") + embedder.calls.clear() + + assert store.delete_document(alpha.id) == alpha + assert embedder.calls == [] + assert store.documents() == [beta] + assert all(hit.doc_id == beta.id for hit in store.retrieve("alpha", k=5)) + assert store.retrieve("beta-specific", k=1)[0].doc_id == beta.id + assert_dense(store) + assert store.add_document("C", "gamma-specific content").id == 2 + + +def test_replace_preserves_document_identity_and_only_embeds_new_chunks(): + embedder = CountingEmbedder() + store = DocumentStore(embedder, min_score=None) + survivor = store.add_document("survivor", "Neptune blue planet") + original = store.add_document("Saturn", "Saturn has distinctive rings") + embedder.calls.clear() + + replaced = store.replace_document( + original.id, "Jupiter", "Jupiter has a giant red storm", max_words=3, overlap=0 + ) + + assert replaced.id == original.id + assert len(embedder.calls) == 1 + assert embedder.calls[0] == ["Jupiter has a", "giant red storm"] + chunks = sorted( + (chunk for chunk in store._chunks.values() if chunk.doc_id == original.id), + key=lambda chunk: chunk.ordinal, + ) + assert [chunk.ordinal for chunk in chunks] == [0, 1] + assert all("Saturn" not in chunk.text for chunk in store._chunks.values()) + assert store.retrieve("Jupiter storm", k=1)[0].doc_id == original.id + assert any(chunk.doc_id == survivor.id for chunk in store._chunks.values()) + assert_dense(store) + + +def test_unknown_lifecycle_targets_are_clear_errors(): + store = DocumentStore(HashedEmbedder(dim=8)) + with pytest.raises(DocumentNotFoundError, match="document 9 was not found"): + store.delete_document(9) + with pytest.raises(DocumentNotFoundError, match="document 9 was not found"): + store.replace_document(9, "missing", "text") + + +def test_delete_last_document_persists_empty_corpus_and_allows_next_insert(tmp_path): + path = tmp_path / "state" + store = DocumentStore(HashedEmbedder(dim=16), state_path=path, min_score=None) + first = store.add_document("only", "only content") + store.delete_document(first.id) + assert store.documents() == [] + assert store.stats()["chunks"] == 0 + assert store.retrieve("only") == [] + + reloaded = DocumentStore(HashedEmbedder(dim=16), state_path=path, min_score=None) + assert reloaded.documents() == [] + assert reloaded.add_document("next", "new content").id == 1 + + +def test_delete_and_replace_survive_restart(tmp_path): + path = tmp_path / "state" + store = DocumentStore(HashedEmbedder(dim=16), state_path=path, min_score=None) + removed = store.add_document("A", "alpha unique") + kept = store.add_document("B", "beta unique") + store.delete_document(removed.id) + store.replace_document(kept.id, "B2", "jupiter unique") + + reloaded = DocumentStore(HashedEmbedder(dim=16), state_path=path, min_score=None) + assert [(doc.id, doc.title) for doc in reloaded.documents()] == [(kept.id, "B2")] + assert reloaded.retrieve("jupiter", k=1)[0].doc_id == kept.id + assert all("alpha" not in chunk.text for chunk in reloaded._chunks.values()) + assert_dense(reloaded) + + +@pytest.mark.parametrize("operation", ["delete", "replace"]) +def test_failed_lifecycle_persistence_leaves_live_and_disk_state_unchanged( + tmp_path, monkeypatch, operation +): + path = tmp_path / "state" + store = DocumentStore(HashedEmbedder(dim=16), state_path=path, min_score=None) + original = store.add_document("original", "saturn original content") + store.add_document("survivor", "neptune survivor content") + before_bytes = path.read_bytes() + before_documents = store.documents() + before_chunks = dict(store._chunks) + monkeypatch.setattr( + store, "_persist", lambda *args: (_ for _ in ()).throw(OSError("disk full")) + ) + + with pytest.raises(OSError, match="disk full"): + if operation == "delete": + store.delete_document(original.id) + else: + store.replace_document(original.id, "replacement", "jupiter replacement") + + assert store.documents() == before_documents + assert store._chunks == before_chunks + assert store.retrieve("saturn", k=1)[0].doc_id == original.id + assert path.read_bytes() == before_bytes + reloaded = DocumentStore(HashedEmbedder(dim=16), state_path=path, min_score=None) + assert reloaded.documents() == before_documents diff --git a/service/tests/test_e2e.py b/service/tests/test_e2e.py index b5b71b3..e73bc5d 100644 --- a/service/tests/test_e2e.py +++ b/service/tests/test_e2e.py @@ -93,6 +93,38 @@ def test_uploads_disabled_do_not_modify_the_index(client): assert after["chunks"] == before["chunks"] +def test_document_lifecycle_endpoints(client, monkeypatch): + monkeypatch.setenv("RAG_UPLOADS_ENABLED", "1") + created = client.post( + "/documents", json={"title": "Saturn", "text": "Saturn ring marker"} + ).json() + replaced = client.put( + f"/documents/{created['id']}", + json={"title": "Jupiter", "text": "Jupiter storm marker"}, + ) + assert replaced.status_code == 200 + assert replaced.json()["id"] == created["id"] + assert replaced.json()["title"] == "Jupiter" + assert client.delete(f"/documents/{created['id']}").json() == replaced.json() + assert created["id"] not in {doc["id"] for doc in client.get("/documents").json()} + assert client.delete(f"/documents/{created['id']}").status_code == 404 + + +def test_lifecycle_endpoints_are_disabled(client): + assert client.delete("/documents/0").status_code == 403 + assert client.put("/documents/0", json={"title": "x", "text": "y"}).status_code == 403 + + +def test_oversized_replacement_body_returns_413(client, monkeypatch): + monkeypatch.setenv("RAG_UPLOADS_ENABLED", "1") + response = client.put( + "/documents/0", + content=b'{"title":"x","text":"' + b"x" * MAX_REQUEST_BYTES + b'"}', + headers={"Content-Type": "application/json"}, + ) + assert response.status_code == 413 + + @pytest.mark.parametrize( "payload", [ @@ -205,6 +237,18 @@ def cors_healthz(): assert allowed.headers["access-control-allow-origin"] == "https://allowed.example" assert "access-control-allow-origin" not in blocked.headers + put = cors_client.options( + "/healthz", + headers={ + "Origin": "https://allowed.example", + "Access-Control-Request-Method": "PUT", + "Access-Control-Request-Headers": "Content-Type", + }, + ) + assert put.status_code == 200 + assert "PUT" in put.headers["access-control-allow-methods"] + assert "DELETE" in put.headers["access-control-allow-methods"] + def test_cors_origin_parser_supports_exact_lists_and_explicit_wildcard(monkeypatch): monkeypatch.setenv(