refactor: add vector backend adapter with Chroma compatibility

paperforge/embedding/__init__.py

@@ -5,6 +5,11 @@ from paperforge.embedding._chroma import (
     get_collection,
     get_vector_db_path,
 )
+from paperforge.embedding.backends import (
+    ChromaBackend,
+    VectorBackend,
+    get_vector_backend,
+)
 from paperforge.embedding.build_state import (
     get_vector_build_state_path,
     mark_vector_build_state,
@@ -17,10 +22,13 @@ from paperforge.embedding.search import retrieve_chunks
 from paperforge.embedding.status import get_embed_status
 
 __all__ = [
+    "ChromaBackend",
+    "VectorBackend",
     "delete_paper_vectors",
     "embed_paper",
     "get_collection",
     "get_embed_status",
+    "get_vector_backend",
     "get_vector_build_state_path",
     "get_vector_db_path",
     "mark_vector_build_state",

paperforge/embedding/backends/__init__.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from paperforge.embedding.backends.base import VectorBackend
+from paperforge.embedding.backends.chroma_backend import ChromaBackend
+
+
+def get_vector_backend(vault: Path) -> VectorBackend:
+    """Return the default vector backend for *vault*."""
+    return ChromaBackend(vault)
+
+
+__all__ = [
+    "ChromaBackend",
+    "VectorBackend",
+    "get_vector_backend",
+]

paperforge/embedding/backends/base.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Protocol
+
+
+class VectorBackend(Protocol):
+    """Adapter protocol for vector database backends.
+
+    The backend wraps all vector DB operations behind a single interface,
+    so the rest of the codebase never depends on a specific vector DB library.
+    """
+
+    def add(
+        self,
+        *,
+        ids: list[str],
+        embeddings: list[list[float]],
+        documents: list[str],
+        metadatas: list[dict],
+    ) -> None: ...
+
+    def query(
+        self, *, query_embedding: list[float], limit: int
+    ) -> list[dict]: ...
+
+    def delete_paper(self, paper_id: str) -> int: ...
+
+    def health(self) -> dict: ...

paperforge/embedding/backends/chroma_backend.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import chromadb
+
+from paperforge.embedding._chroma import get_vector_db_path
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaBackend:
+    """ChromaDB-based vector backend.
+
+    Preserves the existing ``paperforge_fulltext`` collection behaviour.
+    """
+
+    collection_name = "paperforge_fulltext"
+
+    def __init__(self, vault: Path) -> None:
+        db_path = get_vector_db_path(vault)
+        db_path.mkdir(parents=True, exist_ok=True)
+        self.client = chromadb.PersistentClient(path=str(db_path))
+        self.collection = self.client.get_or_create_collection(
+            name=self.collection_name,
+            metadata={"hnsw:space": "cosine"},
+        )
+
+    def add(
+        self,
+        *,
+        ids: list[str],
+        embeddings: list[list[float]],
+        documents: list[str],
+        metadatas: list[dict],
+    ) -> None:
+        try:
+            self.collection.add(
+                ids=ids,
+                embeddings=embeddings,
+                documents=documents,
+                metadatas=metadatas,
+            )
+        except Exception as exc:
+            err = str(exc).lower()
+            if "hnsw" in err or "compaction" in err or "segment" in err:
+                raise RuntimeError(
+                    f"ChromaDB index error (possibly corrupted). "
+                    f"Run 'paperforge embed build --force' to rebuild from scratch. "
+                    f"Original error: {exc}"
+                ) from exc
+            raise
+
+    def query(
+        self, *, query_embedding: list[float], limit: int
+    ) -> list[dict]:
+        raw = self.collection.query(
+            query_embeddings=[query_embedding],
+            n_results=limit,
+            include=["documents", "metadatas", "distances"],
+        )
+        chunks: list[dict] = []
+        for doc, meta, dist in zip(
+            raw["documents"][0],
+            raw["metadatas"][0],
+            raw["distances"][0],
+        ):
+            chunks.append(
+                {
+                    "paper_id": meta["paper_id"],
+                    "section": meta.get("section", "Text"),
+                    "page_number": meta.get("page_number", 1),
+                    "chunk_index": meta.get("chunk_index", 0),
+                    "chunk_text": doc,
+                    "score": round(1.0 - dist, 4),
+                }
+            )
+        return chunks
+
+    def delete_paper(self, paper_id: str) -> int:
+        try:
+            results = self.collection.get(where={"paper_id": paper_id})
+            ids = results.get("ids", [])
+            if ids:
+                self.collection.delete(ids=ids)
+            return len(ids)
+        except Exception:
+            return 0
+
+    def health(self) -> dict:
+        try:
+            chunk_count = self.collection.count()
+            return {"healthy": True, "chunk_count": chunk_count}
+        except Exception as exc:
+            err_lower = str(exc).lower()
+            return {
+                "healthy": False,
+                "chunk_count": 0,
+                "error": str(exc),
+                "corrupted": "hnsw" in err_lower or "corrupt" in err_lower,
+            }

paperforge/embedding/builder.py

@@ -3,15 +3,15 @@ from __future__ import annotations
 import logging
 from pathlib import Path
 
-from paperforge.embedding._chroma import get_collection
+from paperforge.embedding.backends import get_vector_backend
 from paperforge.embedding.providers.openai_compatible import OpenAICompatibleProvider
 
 logger = logging.getLogger(__name__)
 
 
 def embed_paper(vault: Path, zotero_key: str, chunks: list[dict]) -> int:
-    """Embed chunks for one paper using API and insert into ChromaDB. Returns count."""
-    collection = get_collection(vault)
+    """Embed chunks for one paper using API and insert into vector DB. Returns count."""
+    backend = get_vector_backend(vault)
     provider = OpenAICompatibleProvider(vault)
 
     texts = [c["text"] for c in chunks]
@@ -28,20 +28,5 @@ def embed_paper(vault: Path, zotero_key: str, chunks: list[dict]) -> int:
     ]
 
     embeddings = provider.encode(texts)
-    try:
-        collection.add(
-            ids=ids,
-            embeddings=embeddings,
-            documents=texts,
-            metadatas=metadatas,
-        )
-    except Exception as exc:
-        err = str(exc).lower()
-        if "hnsw" in err or "compaction" in err or "segment" in err:
-            raise RuntimeError(
-                f"ChromaDB index error (possibly corrupted). "
-                f"Run 'paperforge embed build --force' to rebuild from scratch. "
-                f"Original error: {exc}"
-            ) from exc
-        raise
+    backend.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas)
     return len(chunks)

paperforge/embedding/search.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import logging
 from pathlib import Path
 
-from paperforge.embedding._chroma import get_collection
+from paperforge.embedding.backends import get_vector_backend
 from paperforge.embedding.providers.openai_compatible import OpenAICompatibleProvider
 
 logger = logging.getLogger(__name__)
@@ -11,29 +11,8 @@ logger = logging.getLogger(__name__)
 
 def retrieve_chunks(vault: Path, query: str, limit: int = 5, expand: bool = True) -> list[dict]:
     """Search chunks via API embedding. Returns list with metadata and similarity scores."""
-    collection = get_collection(vault)
+    backend = get_vector_backend(vault)
     provider = OpenAICompatibleProvider(vault)
     query_embedding = provider.encode_single(query)
 
-    results = collection.query(
-        query_embeddings=[query_embedding],
-        n_results=limit * 3 if expand else limit,
-        include=["documents", "metadatas", "distances"],
-    )
-
-    chunks = []
-    for i, (doc, meta, dist) in enumerate(zip(
-        results["documents"][0],
-        results["metadatas"][0],
-        results["distances"][0],
-    )):
-        chunks.append({
-            "paper_id": meta["paper_id"],
-            "section": meta.get("section", "Text"),
-            "page_number": meta.get("page_number", 1),
-            "chunk_index": meta.get("chunk_index", 0),
-            "chunk_text": doc,
-            "score": round(1.0 - dist, 4),
-        })
-
-    return chunks
+    return backend.query(query_embedding=query_embedding, limit=limit * 3 if expand else limit)

paperforge/embedding/status.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 import logging
 from pathlib import Path
 
-from paperforge.embedding._chroma import get_collection, get_vector_db_path
+from paperforge.embedding._chroma import get_vector_db_path
 from paperforge.embedding._config import get_api_model
+from paperforge.embedding.backends import get_vector_backend
 
 logger = logging.getLogger(__name__)
 
@@ -21,14 +22,12 @@ def get_embed_status(vault: Path) -> dict:
     error = ""
     corrupted = False
     if exists:
-        try:
-            collection = get_collection(vault)
-            chunk_count = collection.count()
-        except Exception as exc:
-            healthy = False
-            error = str(exc)
-            err_lower = str(exc).lower()
-            corrupted = "hnsw" in err_lower or "corrupt" in err_lower
+        backend = get_vector_backend(vault)
+        health = backend.health()
+        healthy = health["healthy"]
+        chunk_count = health.get("chunk_count", 0)
+        error = health.get("error", "")
+        corrupted = health.get("corrupted", False)
 
     model = get_api_model(vault)
 

tests/test_layer4_vector_backend.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from paperforge.embedding.backends import ChromaBackend, get_vector_backend
+
+
+
+
+class TestChromaBackendIdentity:
+    """Identity and naming."""
+
+    def test_chroma_backend_keeps_existing_collection_name(self, tmp_path: Path):
+        backend = ChromaBackend(tmp_path)
+        assert backend.collection_name == "paperforge_fulltext"
+
+    def test_factory_returns_chroma_backend(self, tmp_path: Path):
+        backend = get_vector_backend(tmp_path)
+        assert isinstance(backend, ChromaBackend)
+        assert backend.collection_name == "paperforge_fulltext"
+
+
+class TestChromaBackendAdd:
+    """Adding embeddings."""
+
+    def test_add_passes_through_to_collection(self, chroma_backend: ChromaBackend):
+        ids = ["key_0", "key_1"]
+        embeddings = [[0.1, 0.2], [0.3, 0.4]]
+        documents = ["chunk a", "chunk b"]
+        metadatas = [{"paper_id": "key"}, {"paper_id": "key"}]
+
+        chroma_backend.add(
+            ids=ids,
+            embeddings=embeddings,
+            documents=documents,
+            metadatas=metadatas,
+        )
+
+        chroma_backend.collection.add.assert_called_once_with(
+            ids=ids,
+            embeddings=embeddings,
+            documents=documents,
+            metadatas=metadatas,
+        )
+
+    def test_add_wraps_hnsw_error(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.add.side_effect = RuntimeError("Error loading hnsw index")
+
+        with pytest.raises(RuntimeError, match="ChromaDB index error"):
+            chroma_backend.add(
+                ids=["x"],
+                embeddings=[[0.1]],
+                documents=["t"],
+                metadatas=[{"paper_id": "x"}],
+            )
+
+    def test_add_passes_through_other_errors(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.add.side_effect = ValueError("something else")
+
+        with pytest.raises(ValueError, match="something else"):
+            chroma_backend.add(
+                ids=["x"],
+                embeddings=[[0.1]],
+                documents=["t"],
+                metadatas=[{"paper_id": "x"}],
+            )
+
+
+class TestChromaBackendQuery:
+    """Querying embeddings."""
+
+    def test_query_returns_formatted_chunks(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.query.return_value = {
+            "documents": [["chunk text"]],
+            "metadatas": [[{"paper_id": "abc", "section": "Methods", "page_number": 3, "chunk_index": 1}]],
+            "distances": [[0.15]],
+        }
+
+        results = chroma_backend.query(query_embedding=[0.1, 0.2], limit=5)
+
+        assert len(results) == 1
+        assert results[0]["paper_id"] == "abc"
+        assert results[0]["section"] == "Methods"
+        assert results[0]["page_number"] == 3
+        assert results[0]["chunk_index"] == 1
+        assert results[0]["chunk_text"] == "chunk text"
+        assert results[0]["score"] == 0.85  # 1.0 - 0.15
+
+    def test_query_with_missing_metadata_fields(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.query.return_value = {
+            "documents": [["text"]],
+            "metadatas": [[{"paper_id": "abc"}]],
+            "distances": [[0.3]],
+        }
+
+        results = chroma_backend.query(query_embedding=[0.1], limit=5)
+
+        assert results[0]["section"] == "Text"
+        assert results[0]["page_number"] == 1
+        assert results[0]["chunk_index"] == 0
+
+    def test_query_passes_limit(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.query.return_value = {
+            "documents": [[]],
+            "metadatas": [[]],
+            "distances": [[]],
+        }
+
+        chroma_backend.query(query_embedding=[0.1], limit=10)
+
+        chroma_backend.collection.query.assert_called_once()
+        assert chroma_backend.collection.query.call_args[1]["n_results"] == 10
+
+
+class TestChromaBackendDelete:
+    """Deleting paper vectors."""
+
+    def test_delete_paper_deletes_found_ids(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.get.return_value = {"ids": ["abc_0", "abc_1"]}
+
+        count = chroma_backend.delete_paper("abc")
+
+        assert count == 2
+        chroma_backend.collection.delete.assert_called_once_with(ids=["abc_0", "abc_1"])
+
+    def test_delete_paper_no_ids(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.get.return_value = {"ids": []}
+
+        count = chroma_backend.delete_paper("abc")
+
+        assert count == 0
+        chroma_backend.collection.delete.assert_not_called()
+
+    def test_delete_paper_handles_exception(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.get.side_effect = RuntimeError("whoops")
+
+        count = chroma_backend.delete_paper("abc")
+
+        assert count == 0
+
+
+class TestChromaBackendHealth:
+    """Health check."""
+
+    def test_health_returns_healthy(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.count.return_value = 42
+
+        result = chroma_backend.health()
+
+        assert result["healthy"] is True
+        assert result["chunk_count"] == 42
+
+    def test_health_returns_unhealthy(self, chroma_backend: ChromaBackend):
+        chroma_backend.collection.count.side_effect = RuntimeError("corrupt index")
+
+        result = chroma_backend.health()
+
+        assert result["healthy"] is False
+        assert result["chunk_count"] == 0
+        assert "corrupt" in result.get("error", "")
+        assert result.get("corrupted") is True
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def chroma_backend(tmp_path: Path) -> ChromaBackend:
+    """Return a ChromaBackend whose client and collection are fully mocked."""
+    mock_client = MagicMock()
+    mock_collection = MagicMock()
+    mock_client.get_or_create_collection.return_value = mock_collection
+
+    with patch("paperforge.embedding.backends.chroma_backend.chromadb.PersistentClient", return_value=mock_client):
+        backend = ChromaBackend(tmp_path)
+
+    # Attach the mock collection for assertion convenience
+    backend.collection = mock_collection
+    backend.client = mock_client
+    return backend

tests/unit/worker/test_vector_db.py

@@ -5,15 +5,25 @@ from unittest.mock import Mock, patch
 from paperforge.embedding.status import get_embed_status
 
 
+def _mock_backend(health_result: dict) -> Mock:
+    backend = Mock()
+    backend.health.return_value = health_result
+    return backend
+
+
 def test_get_embed_status_reports_corruption_when_count_fails(tmp_path):
     vault = tmp_path / "vault"
     vectors_dir = vault / "System" / "PaperForge" / "indexes" / "vectors"
     vectors_dir.mkdir(parents=True)
 
-    mock_collection = Mock()
-    mock_collection.count.side_effect = RuntimeError("Error loading hnsw index")
+    mock_backend = _mock_backend({
+        "healthy": False,
+        "chunk_count": 0,
+        "error": "Error loading hnsw index",
+        "corrupted": True,
+    })
 
-    with patch("paperforge.embedding.status.get_collection", return_value=mock_collection):
+    with patch("paperforge.embedding.status.get_vector_backend", return_value=mock_backend):
         status = get_embed_status(vault)
 
     assert status["db_exists"] is True
@@ -32,10 +42,12 @@ def test_get_embed_status_uses_indexes_vectors_path_from_config(tmp_path):
     vectors_dir = vault / "System" / "PaperForge" / "indexes" / "vectors"
     vectors_dir.mkdir(parents=True)
 
-    mock_collection = Mock()
-    mock_collection.count.return_value = 12
+    mock_backend = _mock_backend({
+        "healthy": True,
+        "chunk_count": 12,
+    })
 
-    with patch("paperforge.embedding.status.get_collection", return_value=mock_collection):
+    with patch("paperforge.embedding.status.get_vector_backend", return_value=mock_backend):
         status = get_embed_status(vault)
 
     assert status["db_exists"] is True