diff --git a/blueprints/rag/logic.py b/blueprints/rag/logic.py index 94ee7ae..665f97b 100644 --- a/blueprints/rag/logic.py +++ b/blueprints/rag/logic.py @@ -1,4 +1,6 @@ +import asyncio import datetime +import json import logging import os import re @@ -325,6 +327,57 @@ def _load_obsidian_index_cache() -> dict[str, float]: return {row[0]: row[1] for row in result if row[0] is not None} +def _walk_vault_mtimes(obsidian_service: ObsidianService) -> dict[str, float]: + """Walk the vault and stat every file. Blocking; run via asyncio.to_thread.""" + return { + str(md_path): md_path.stat().st_mtime + for md_path in obsidian_service.walk_vault() + } + + +def _read_obsidian_document( + obsidian_service: ObsidianService, filepath: str, now: float +) -> Document | None: + """Read and parse one vault file into a Document, or None on error.""" + try: + with open(filepath, "r", encoding="utf-8") as f: + content = f.read() + parsed = obsidian_service.parse_markdown(content, filepath) + metadata = { + "source": "obsidian", + "filepath": parsed["filepath"], + "folder": str(Path(parsed["filepath"]).parent) + if parsed["filepath"] + else "", + "tags": parsed["tags"], + "created_at": parsed["metadata"].get("created_at"), + "indexed_at": now, + **{ + k: v + for k, v in parsed["metadata"].items() + if k not in ["created_at", "created_by"] + }, + } + return Document( + page_content=parsed["content"], metadata=_make_serializable(metadata) + ) + except Exception as e: + logger.warning(f"Error reading {filepath}: {e}") + return None + + +def _read_obsidian_documents( + obsidian_service: ObsidianService, filepaths: list[str], now: float +) -> list[Document]: + """Read and parse a batch of vault files. Blocking; run via asyncio.to_thread.""" + documents = [] + for filepath in filepaths: + document = _read_obsidian_document(obsidian_service, filepath, now) + if document is not None: + documents.append(document) + return documents + + async def sync_obsidian_documents() -> dict[str, int]: """Incrementally sync Obsidian documents to pgvector. @@ -340,69 +393,52 @@ async def sync_obsidian_documents() -> dict[str, int]: # Load cache from DB on first run if not _obsidian_index_cache: - _obsidian_index_cache = _load_obsidian_index_cache() + _obsidian_index_cache = await asyncio.to_thread(_load_obsidian_index_cache) # Build map of current vault files -> mtime - vault_files: dict[str, float] = {} - for md_path in obsidian_service.walk_vault(): - vault_files[str(md_path)] = md_path.stat().st_mtime + vault_files = await asyncio.to_thread(_walk_vault_mtimes, obsidian_service) added = 0 updated = 0 - deleted = 0 # Find files to add or update files_to_index: list[str] = [] + stale_filepaths: list[str] = [] for filepath, mtime in vault_files.items(): indexed_at = _obsidian_index_cache.get(filepath) if indexed_at is None: files_to_index.append(filepath) added += 1 elif mtime > indexed_at: - # Delete old chunks first - delete_documents_by_metadata("filepath", filepath) + stale_filepaths.append(filepath) files_to_index.append(filepath) updated += 1 - # Find deleted files (in cache but not on disk) - for filepath in list(_obsidian_index_cache): - if filepath not in vault_files: - delete_documents_by_metadata("filepath", filepath) - del _obsidian_index_cache[filepath] - deleted += 1 + # Find deleted files (in cache but not on disk). Don't drop them from the + # cache yet: if the delete below raises, we want the next sync to still + # know about them and retry, rather than leaving their chunks searchable + # forever with no record that they need cleaning up. + removed_filepaths = [fp for fp in _obsidian_index_cache if fp not in vault_files] + stale_filepaths.extend(removed_filepaths) + deleted = len(removed_filepaths) + + # Delete old chunks for updated/removed files before re-indexing, in one + # transaction instead of a round-trip per file. Left uncaught: a failure + # here must stop the sync so the cache isn't marked clean for work that + # didn't happen. + if stale_filepaths: + await asyncio.to_thread(delete_documents_by_filepaths, stale_filepaths) + + # Only now that the delete has succeeded is it safe to forget these files. + for filepath in removed_filepaths: + del _obsidian_index_cache[filepath] # Index new/changed files if files_to_index: now = time.time() - documents = [] - for filepath in files_to_index: - try: - with open(filepath, "r", encoding="utf-8") as f: - content = f.read() - parsed = obsidian_service.parse_markdown(content, filepath) - metadata = { - "source": "obsidian", - "filepath": parsed["filepath"], - "folder": str(Path(parsed["filepath"]).parent) - if parsed["filepath"] - else "", - "tags": parsed["tags"], - "created_at": parsed["metadata"].get("created_at"), - "indexed_at": now, - **{ - k: v - for k, v in parsed["metadata"].items() - if k not in ["created_at", "created_by"] - }, - } - document = Document( - page_content=parsed["content"], - metadata=_make_serializable(metadata), - ) - documents.append(document) - except Exception as e: - logger.warning(f"Error reading {filepath}: {e}") - continue + documents = await asyncio.to_thread( + _read_obsidian_documents, obsidian_service, files_to_index, now + ) if documents: splits = [] @@ -507,6 +543,27 @@ def delete_all_documents(): conn.commit() +def _delete_by_containment(conn, collection_id, filter_dict: dict) -> None: + """Delete rows whose cmetadata contains filter_dict. + + Uses JSONB containment (`@>`) rather than `cmetadata->>key = value` so the + delete can use the GIN index langchain_postgres creates on cmetadata + (ix_cmetadata_gin, jsonb_path_ops); `->>` text extraction can't use it. + + Unlike `->>`, which coerces to text before comparing, `@>` matches on the + stored JSON type: the string "123" does not match the number 123. Callers + must pass values with the same type as what was stored (today's only + caller, filepath/source strings, is fine). + """ + conn.execute( + text( + "DELETE FROM langchain_pg_embedding " + "WHERE collection_id = :cid AND cmetadata @> CAST(:filter AS jsonb)" + ), + {"cid": collection_id, "filter": json.dumps(filter_dict)}, + ) + + def delete_documents_by_metadata(key: str, value: str): """Delete documents matching a metadata key/value pair.""" collection_id = _get_collection_id() @@ -514,13 +571,25 @@ def delete_documents_by_metadata(key: str, value: str): return engine = _get_engine() with engine.connect() as conn: - conn.execute( - text( - "DELETE FROM langchain_pg_embedding " - "WHERE collection_id = :cid AND cmetadata->>:key = :value" - ), - {"cid": collection_id, "key": key, "value": value}, - ) + _delete_by_containment(conn, collection_id, {key: value}) + conn.commit() + + +def delete_documents_by_filepaths(filepaths: list[str]) -> None: + """Delete all chunks for several obsidian filepaths in one transaction. + + One containment delete per path, committed together, so a sync pass with + many stale files doesn't pay a commit round-trip per file. + """ + if not filepaths: + return + collection_id = _get_collection_id() + if not collection_id: + return + engine = _get_engine() + with engine.connect() as conn: + for filepath in filepaths: + _delete_by_containment(conn, collection_id, {"filepath": filepath}) conn.commit() diff --git a/tests/unit/test_obsidian_sync.py b/tests/unit/test_obsidian_sync.py new file mode 100644 index 0000000..57b6411 --- /dev/null +++ b/tests/unit/test_obsidian_sync.py @@ -0,0 +1,188 @@ +"""Tests for the incremental Obsidian sync: off-loop execution and delete SQL.""" + +import json +import os +import threading +from unittest.mock import patch + +import pytest + +from blueprints.rag import logic + + +class FakeConnection: + def __init__(self, store): + self.store = store + + def execute(self, stmt, params=None): + self.store.append((str(stmt), params or {})) + + def commit(self): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class FakeEngine: + def __init__(self): + self.executed: list[tuple[str, dict]] = [] + + def connect(self): + return FakeConnection(self.executed) + + +class FakeVectorStore: + def __init__(self): + self.calls: list[list[str]] = [] + + async def aadd_documents(self, documents): + self.calls.append([d.page_content for d in documents]) + + +@pytest.fixture +def vault_dir(tmp_path): + (tmp_path / "new.md").write_text("# New\nNew note body") + (tmp_path / "updated.md").write_text("# Updated\nUpdated note body") + with patch.dict(os.environ, {"OBSIDIAN_VAULT_PATH": str(tmp_path)}): + yield tmp_path + + +@pytest.fixture +async def synced(monkeypatch, vault_dir): + """Run a sync with a fake vault: one new file, one updated file, one + file that's been deleted off disk since the last sync.""" + updated_path = str(vault_dir / "updated.md") + missing_path = str(vault_dir / "gone.md") + + monkeypatch.setattr( + logic, "_obsidian_index_cache", {updated_path: 1.0, missing_path: 1.0} + ) + + engine = FakeEngine() + monkeypatch.setattr(logic, "_get_engine", lambda: engine) + monkeypatch.setattr(logic, "_get_collection_id", lambda: "collection-uuid") + + def _unexpected_cache_load(): + raise AssertionError("cache load should be skipped when already populated") + + monkeypatch.setattr(logic, "_load_obsidian_index_cache", _unexpected_cache_load) + + store = FakeVectorStore() + monkeypatch.setattr(logic, "_get_vector_store", lambda: store) + + threads: dict[str, threading.Thread] = {} + + real_walk = logic._walk_vault_mtimes + + def recording_walk(service): + threads["walk"] = threading.current_thread() + return real_walk(service) + + monkeypatch.setattr(logic, "_walk_vault_mtimes", recording_walk) + + real_read = logic._read_obsidian_documents + + def recording_read(service, filepaths, now): + threads["read"] = threading.current_thread() + return real_read(service, filepaths, now) + + monkeypatch.setattr(logic, "_read_obsidian_documents", recording_read) + + real_delete = logic.delete_documents_by_filepaths + + def recording_delete(filepaths): + threads["delete"] = threading.current_thread() + return real_delete(filepaths) + + monkeypatch.setattr(logic, "delete_documents_by_filepaths", recording_delete) + + result = await logic.sync_obsidian_documents() + return result, threads, engine, store, updated_path, missing_path + + +class TestSyncOffLoop: + async def test_counts_match_vault_changes(self, synced): + result, *_ = synced + assert result == {"added": 1, "updated": 1, "deleted": 1} + + async def test_walk_runs_off_the_main_thread(self, synced): + _, threads, *_ = synced + assert threads["walk"] is not threading.main_thread() + + async def test_reads_run_off_the_main_thread(self, synced): + _, threads, *_ = synced + assert threads["read"] is not threading.main_thread() + + async def test_deletes_run_off_the_main_thread(self, synced): + _, threads, *_ = synced + assert threads["delete"] is not threading.main_thread() + + async def test_embeds_new_and_updated_content(self, synced): + _, _, _, store, *_ = synced + embedded = {text for call in store.calls for text in call} + assert any("New note body" in t for t in embedded) + assert any("Updated note body" in t for t in embedded) + + async def test_cache_reflects_add_update_delete(self, synced): + result, _, _, _, updated_path, missing_path = synced + assert missing_path not in logic._obsidian_index_cache + assert updated_path in logic._obsidian_index_cache + + async def test_deletes_use_jsonb_containment_with_bound_params(self, synced): + _, _, engine, *_ = synced + deletes = [(sql, params) for sql, params in engine.executed if "DELETE" in sql] + assert deletes, "expected at least one delete statement" + for sql, params in deletes: + assert "cmetadata @> CAST(:filter AS jsonb)" in sql + assert "->>:" not in sql # not the old, unindexed text-extraction filter + assert "filter" in params + parsed_filter = json.loads(params["filter"]) + assert "filepath" in parsed_filter + + +async def test_delete_failure_keeps_deleted_path_in_cache_and_propagates( + monkeypatch, vault_dir +): + """If the batched delete raises, the sync must not have already forgotten + the deleted file: otherwise the next sync never retries the delete and + its chunks stay searchable forever with no record that they're stale.""" + updated_path = str(vault_dir / "updated.md") + missing_path = str(vault_dir / "gone.md") + + monkeypatch.setattr( + logic, "_obsidian_index_cache", {updated_path: 1.0, missing_path: 1.0} + ) + + def _raise_on_delete(filepaths): + raise RuntimeError("db unavailable") + + monkeypatch.setattr(logic, "delete_documents_by_filepaths", _raise_on_delete) + + store = FakeVectorStore() + monkeypatch.setattr(logic, "_get_vector_store", lambda: store) + + with pytest.raises(RuntimeError, match="db unavailable"): + await logic.sync_obsidian_documents() + + assert missing_path in logic._obsidian_index_cache + # The updated file's old chunks were never actually deleted, so its cache + # entry must not be bumped to look freshly indexed. + assert logic._obsidian_index_cache[updated_path] == 1.0 + assert store.calls == [] + + +async def test_delete_documents_by_metadata_uses_containment(monkeypatch): + engine = FakeEngine() + monkeypatch.setattr(logic, "_get_engine", lambda: engine) + monkeypatch.setattr(logic, "_get_collection_id", lambda: "collection-uuid") + + logic.delete_documents_by_metadata("source", "obsidian") + + assert len(engine.executed) == 1 + sql, params = engine.executed[0] + assert "cmetadata @> CAST(:filter AS jsonb)" in sql + assert json.loads(params["filter"]) == {"source": "obsidian"}