/qdrant/
+
+No additional sidecar files are needed for the vector data itself (unlike local).
+The embed_id_to_extra, doc_ids_to_file_names, metadata, and index_config sidecars
+are still managed by ColPaliModel directly.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, ClassVar, Dict, List, Optional
+
+import torch
+
+from .base import (
+ MultiVectorQuery,
+ SearchHit,
+ StoredPoint,
+ VectorStore,
+)
+
+logger = logging.getLogger(__name__)
+
+try:
+ from qdrant_client import QdrantClient
+ from qdrant_client.models import (
+ Distance,
+ FieldCondition,
+ Filter,
+ MatchValue,
+ MultiVectorComparator,
+ MultiVectorConfig,
+ PointStruct,
+ VectorParams,
+ )
+ _QDRANT_AVAILABLE = True
+except ImportError:
+ _QDRANT_AVAILABLE = False
+
+
+def _require_qdrant() -> None:
+ if not _QDRANT_AVAILABLE:
+ raise RuntimeError(
+ "The Qdrant storage backend requires the qdrant-client package.\n"
+ "Install it with: pip install \"foretrieval[qdrant]\"\n"
+ "or: uv add foretrieval --extra qdrant"
+ )
+
+
+class QdrantVectorStore(VectorStore):
+ """Embedded Qdrant vector store with native multi-vector MAX_SIM scoring."""
+
+ backend_name: ClassVar[str] = "qdrant"
+ supports_multivector_native: ClassVar[bool] = True
+
+ def __init__(self) -> None:
+ self._client: Optional["QdrantClient"] = None
+ self._collection_name: Optional[str] = None
+ self._index_root: Optional[Path] = None
+
+ # ------------------------------------------------------------------
+ # Lifecycle
+ # ------------------------------------------------------------------
+
+ def open(
+ self,
+ index_name: str,
+ index_root: Path,
+ *,
+ create: bool,
+ dim: Optional[int] = None,
+ ) -> None:
+ _require_qdrant()
+ self._collection_name = index_name
+ self._index_root = Path(index_root)
+ qdrant_path = self._index_root / index_name / "qdrant"
+ qdrant_path.mkdir(parents=True, exist_ok=True)
+ self._client = QdrantClient(path=str(qdrant_path))
+
+ if create and dim is not None and not self._client.collection_exists(index_name):
+ self.create_collection(dim)
+
+ def close(self) -> None:
+ self._client = None
+
+ # ------------------------------------------------------------------
+ # Collection management
+ # ------------------------------------------------------------------
+
+ def collection_exists(self) -> bool:
+ if self._client is None or self._collection_name is None:
+ return False
+ return self._client.collection_exists(self._collection_name)
+
+ def create_collection(self, dim: int) -> None:
+ if self._client is None or self._collection_name is None:
+ raise RuntimeError("QdrantVectorStore.open() must be called first.")
+ if self._client.collection_exists(self._collection_name):
+ return
+ self._client.create_collection(
+ collection_name=self._collection_name,
+ vectors_config=VectorParams(
+ size=dim,
+ distance=Distance.COSINE,
+ multivector_config=MultiVectorConfig(
+ comparator=MultiVectorComparator.MAX_SIM,
+ ),
+ ),
+ )
+
+ # ------------------------------------------------------------------
+ # Write
+ # ------------------------------------------------------------------
+
+ def upsert(self, points: List[StoredPoint]) -> None:
+ if self._client is None or self._collection_name is None:
+ raise RuntimeError("QdrantVectorStore.open() must be called first.")
+ qdrant_points = [
+ PointStruct(
+ id=sp.point_id,
+ vector=sp.vector.float().numpy().tolist(),
+ payload=sp.payload,
+ )
+ for sp in points
+ ]
+ self._client.upsert(
+ collection_name=self._collection_name,
+ points=qdrant_points,
+ )
+
+ def point_exists(self, point_id: int) -> bool:
+ if self._client is None or self._collection_name is None:
+ return False
+ if not self._client.collection_exists(self._collection_name):
+ return False
+ found = self._client.retrieve(
+ collection_name=self._collection_name,
+ ids=[point_id],
+ with_payload=False,
+ with_vectors=False,
+ )
+ return bool(found)
+
+ # ------------------------------------------------------------------
+ # Read
+ # ------------------------------------------------------------------
+
+ def search(
+ self,
+ query: MultiVectorQuery,
+ k: int,
+ ) -> List[SearchHit]:
+ if self._client is None or self._collection_name is None:
+ raise RuntimeError("QdrantVectorStore.open() must be called first.")
+
+ qfilter = self._build_filter(query.filter_metadata)
+
+ response = self._client.query_points(
+ collection_name=self._collection_name,
+ query=query.vectors.float().numpy().tolist(),
+ query_filter=qfilter,
+ limit=k,
+ with_payload=True,
+ with_vectors=False,
+ )
+
+ points = response.points if hasattr(response, "points") else response
+
+ return [
+ SearchHit(
+ point_id=int(p.id),
+ score=float(p.score),
+ payload=p.payload or {},
+ )
+ for p in points
+ ]
+
+ def fetch_vector(self, point_id: int) -> Optional[torch.Tensor]:
+ if self._client is None or self._collection_name is None:
+ return None
+ retrieved = self._client.retrieve(
+ collection_name=self._collection_name,
+ ids=[point_id],
+ with_payload=False,
+ with_vectors=True,
+ )
+ if not retrieved:
+ return None
+ return torch.tensor(retrieved[0].vector)
+
+ # ------------------------------------------------------------------
+ # Filter helper
+ # ------------------------------------------------------------------
+
+ def _build_filter(
+ self, filter_metadata: Optional[Dict[str, Any]]
+ ) -> Optional["Filter"]:
+ if not filter_metadata:
+ return None
+ must = [
+ FieldCondition(
+ key=f"metadata.{k}",
+ match=MatchValue(value=v),
+ )
+ for k, v in filter_metadata.items()
+ ]
+ return Filter(must=must)
+
+ # ------------------------------------------------------------------
+ # Client accessor (for testing / ColPaliModel compatibility)
+ # ------------------------------------------------------------------
+
+ @property
+ def client(self) -> Optional["QdrantClient"]:
+ return self._client
+
+ @property
+ def collection_name(self) -> Optional[str]:
+ return self._collection_name
diff --git a/foretrieval/vector_store/remote.py b/foretrieval/vector_store/remote.py
new file mode 100644
index 0000000..5d8b0bb
--- /dev/null
+++ b/foretrieval/vector_store/remote.py
@@ -0,0 +1,196 @@
+"""RemoteVectorStore — VectorStore client that delegates to a remote server.
+
+This is the client-side VectorStore implementation. It holds a
+VectorDBServerClient and implements the VectorStore ABC by forwarding every
+call to the corresponding HTTP endpoint on the remote server.
+
+``export_sidecar`` and ``load_sidecar`` are no-ops because the data lives on
+the server, not on the client filesystem.
+
+Usage:
+ From make_vector_store() factory (preferred):
+
+ vs = make_vector_store(
+ "remote",
+ {
+ "url": "http://gpu-server:18000",
+ "backend": "qdrant", # server-side backend
+ "api_key": "secret", # optional
+ }
+ )
+ vs.open("my_index", Path(".foretrieval"), create=True, dim=128)
+
+ Or directly:
+
+ from foretrieval.vector_db_server import VectorDBServerClient, VectorDBServerConfig
+ from foretrieval.vector_store.remote import RemoteVectorStore
+
+ cfg = VectorDBServerConfig(url="http://gpu-server:18000", backend="qdrant")
+ client = VectorDBServerClient(cfg)
+ vs = RemoteVectorStore(client, backend="qdrant")
+ vs.open("my_index", Path("."), create=True, dim=128)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, ClassVar, Dict, List, Optional
+
+import torch
+
+from .base import MultiVectorQuery, SearchHit, StoredPoint, VectorStore
+
+
+class RemoteVectorStore(VectorStore):
+ """VectorStore that delegates to a remote FORetrieval vector-DB server.
+
+ Parameters
+ ----------
+ client:
+ VectorDBServerClient connected to the target server.
+ backend:
+ Server-side storage backend (``"local"``, ``"qdrant"``, or
+ ``"milvus"``). Used when creating a new collection.
+ storage_config:
+ Optional backend-specific config forwarded to the server
+ (e.g. ``{"candidate_limit": 128}`` for Milvus).
+ """
+
+ backend_name: ClassVar[str] = "remote"
+ supports_multivector_native: ClassVar[bool] = True
+ # (informational — depends on the server-side backend, but we report True
+ # since qdrant/local are both exact; milvus is approximate but this
+ # attribute is only used for informational logging)
+
+ def __init__(
+ self,
+ client: Any, # VectorDBServerClient — avoid circular import at module level
+ backend: str = "qdrant",
+ storage_config: Optional[Dict[str, Any]] = None,
+ ) -> None:
+ self._client = client
+ self._backend = backend
+ self._storage_config = storage_config
+
+ # Set by open()
+ self._index_name: Optional[str] = None
+ self._opened: bool = False
+
+ # ------------------------------------------------------------------
+ # Lifecycle
+ # ------------------------------------------------------------------
+
+ def open(
+ self,
+ index_name: str,
+ index_root: Path,
+ *,
+ create: bool,
+ dim: Optional[int] = None,
+ ) -> None:
+ """Open or create the remote collection.
+
+ ``index_root`` is ignored — data lives on the server.
+ """
+ self._index_name = index_name
+ self._client.open_collection(
+ index_name,
+ self._backend,
+ create=create,
+ dim=dim,
+ storage_config=self._storage_config,
+ )
+ self._opened = True
+
+ def close(self) -> None:
+ """Close the underlying HTTP client."""
+ self._client.close()
+ self._opened = False
+
+ # ------------------------------------------------------------------
+ # Collection management
+ # ------------------------------------------------------------------
+
+ def collection_exists(self) -> bool:
+ if self._index_name is None:
+ return False
+ return self._client.collection_exists(self._index_name)
+
+ def create_collection(self, dim: int) -> None:
+ if self._index_name is None:
+ raise RuntimeError("Call open() before create_collection()")
+ self._client.create_collection(
+ self._index_name, self._backend, dim, self._storage_config
+ )
+
+ # ------------------------------------------------------------------
+ # Write
+ # ------------------------------------------------------------------
+
+ def upsert(self, points: List[StoredPoint]) -> None:
+ if self._index_name is None:
+ raise RuntimeError("Call open() before upsert()")
+ self._client.upsert(self._index_name, points)
+
+ def point_exists(self, point_id: int) -> bool:
+ if self._index_name is None:
+ return False
+ return self._client.point_exists(self._index_name, point_id)
+
+ # ------------------------------------------------------------------
+ # Read
+ # ------------------------------------------------------------------
+
+ def search(
+ self,
+ query: MultiVectorQuery,
+ k: int,
+ ) -> List[SearchHit]:
+ if self._index_name is None:
+ raise RuntimeError("Call open() before search()")
+ return self._client.search(self._index_name, query, k)
+
+ def fetch_vector(self, point_id: int) -> Optional[torch.Tensor]:
+ if self._index_name is None:
+ return None
+ return self._client.fetch_vector(self._index_name, point_id)
+
+ # ------------------------------------------------------------------
+ # Persistence helpers — no-ops for remote store
+ # ------------------------------------------------------------------
+
+ def export_sidecar(self, index_path: Path) -> None:
+ """No-op — vector data lives on the server."""
+
+ def load_sidecar(self, index_path: Path) -> None:
+ """No-op — vector data lives on the server."""
+
+ # ------------------------------------------------------------------
+ # Index-level bookkeeping — round-tripped to the server so the client
+ # needs no local index directory.
+ # ------------------------------------------------------------------
+
+ def supports_remote_bookkeeping(self) -> bool:
+ return True
+
+ def export_bookkeeping(self, blob: Dict[str, Any]) -> None:
+ if self._index_name is None:
+ raise RuntimeError("Call open() before export_bookkeeping()")
+ self._client.put_bookkeeping(self._index_name, blob)
+
+ def load_bookkeeping(self) -> Optional[Dict[str, Any]]:
+ if self._index_name is None:
+ raise RuntimeError("Call open() before load_bookkeeping()")
+ return self._client.get_bookkeeping(self._index_name)
+
+ # ------------------------------------------------------------------
+ # Internal accessors (for ColPaliModel compatibility checks)
+ # ------------------------------------------------------------------
+
+ @property
+ def is_opened(self) -> bool:
+ return self._opened
+
+ @property
+ def server_backend(self) -> str:
+ return self._backend
diff --git a/outside_world/docs/attention.pdf b/outside_world/docs/attention.pdf
deleted file mode 100644
index 97d7c51..0000000
Binary files a/outside_world/docs/attention.pdf and /dev/null differ
diff --git a/outside_world/docs/attention_with_a_mustache.pdf b/outside_world/docs/attention_with_a_mustache.pdf
deleted file mode 100644
index 97d7c51..0000000
Binary files a/outside_world/docs/attention_with_a_mustache.pdf and /dev/null differ
diff --git a/pyproject.toml b/pyproject.toml
index bf1d624..695cb7b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,14 +3,14 @@ requires = ["setuptools"]
build-backend = "setuptools.build_meta"
[tool.setuptools.packages.find]
-include = ["byaldi*"]
+include = ["foretrieval*"]
[project]
-name = "Byaldi"
-version = "0.0.7"
+name = "FORetrieval"
+version = "2026.6.3"
description = "Use late-interaction multi-modal models such as ColPali in just a few lines of code."
readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.12"
license = { file = "LICENSE" }
keywords = [
"reranking",
@@ -21,34 +21,61 @@ keywords = [
"colbert",
"multi-modal",
]
-authors = [{ name = "Ben Clavié", email = "bc@answer.ai" }]
-maintainers = [
- { name = "Ben Clavié", email = "bc@answer.ai" },
- { name = "Tony Wu", email = "tony.wu@illuin.tech" },
-]
+authors = [{ name = "The FOR team", email = "for@irt-saintexupery.com" }]
dependencies = [
- "colpali-engine>=0.3.4,<0.4.0",
+ "colpali-engine>=0.3.15,<0.4.0",
+ "httpx>=0.27.0",
+ "docx2pdf>=0.1.8",
+ "langdetect>=1.0.9",
"ml-dtypes",
- "mteb==1.6.35",
- "ninja",
- "pdf2image",
+ "pydantic-ai>=1.4.0,<=1.107.0",
+ "pydantic-ai-slim[openai]>=1.8.0",
+ "pypdf>=6.1.3",
"srsly",
- "torch",
- "transformers>=4.42.0",
+ "torch>=2.7.1",
+ "transformers>=4.42.0,<5.6.0",
+ "pdf2image>=1.17.0",
]
[project.optional-dependencies]
-dev = ["pytest>=7.4.0", "ruff>=0.1.9"]
-server = ["uvicorn", "fastapi"]
+# Qdrant vector-store backend (recommended for large indexes)
+qdrant = ["qdrant-client>=1.17.1"]
+# Milvus vector-store backend (Milvus Lite for local dev, full Milvus for prod)
+milvus = ["pymilvus>=2.4.0"]
+# Docling-based PDF chunking ingestion pipeline
+docling = ["docling>=2.76.0"]
+embedding_server = ["paramiko>=3.0"]
+# Remote vector-DB server (server + auto-deploy; client needs only httpx which is a core dep)
+vector_db_server = ["fastapi>=0.115", "uvicorn>=0.30", "python-multipart>=0.0.9", "paramiko>=3.0", "matplotlib>=3.7.0"]
+# 4-bit / 8-bit quantization for local ColPaliModel (requires CUDA)
+quantization = ["bitsandbytes>=0.42.0"]
+dev = ["pytest>=7.4.0", "pytest-mock>=3.0", "ruff>=0.1.9"]
langchain = ["langchain-core"]
+extra_converters = [
+ "docx2pdf>=0.1.8; sys_platform == \"win32\"",
+ "reportlab>=4.4.4",
+ "python-docx>=1.2.0",
+ "python-pptx>=1.0.2",
+]
[project.urls]
-"Homepage" = "https://github.com/answerdotai/byaldi"
+"Homepage" = "https://github.com/FOR-sight-ai/FORetrieval"
+
+[project.scripts]
+foretrieval-db-server = "foretrieval.vector_db_server.server_main:main"
+
+[dependency-groups]
+dev = [
+ "matplotlib>=3.10.8",
+]
[tool.pytest.ini_options]
filterwarnings = ["ignore::Warning"]
-markers = ["slow: marks test as slow"]
+markers = [
+ "slow: marks test as slow (GPU-dependent or expensive)",
+ "integration: marks test as requiring a live API key",
+]
testpaths = ["tests"]
[tool.ruff]
@@ -83,7 +110,7 @@ exclude = [
line-length = 88
output-format = "grouped"
-target-version = "py39"
+target-version = "py312"
[tool.ruff.lint]
select = [
diff --git a/sample_data/sample_bmp.bmp b/sample_data/sample_bmp.bmp
new file mode 100644
index 0000000..21d9335
Binary files /dev/null and b/sample_data/sample_bmp.bmp differ
diff --git a/sample_data/sample_csv.csv b/sample_data/sample_csv.csv
new file mode 100644
index 0000000..2eeba7b
--- /dev/null
+++ b/sample_data/sample_csv.csv
@@ -0,0 +1 @@
+This is a sample .csv file, The answer you're looking for: 484696
\ No newline at end of file
diff --git a/sample_data/sample_csv.pdf b/sample_data/sample_csv.pdf
new file mode 100644
index 0000000..4c0920f
Binary files /dev/null and b/sample_data/sample_csv.pdf differ
diff --git a/sample_data/sample_doc.doc b/sample_data/sample_doc.doc
new file mode 100644
index 0000000..e52a1f4
Binary files /dev/null and b/sample_data/sample_doc.doc differ
diff --git a/sample_data/sample_doc.pdf b/sample_data/sample_doc.pdf
new file mode 100644
index 0000000..69efbeb
Binary files /dev/null and b/sample_data/sample_doc.pdf differ
diff --git a/sample_data/sample_docx.docx b/sample_data/sample_docx.docx
new file mode 100644
index 0000000..8886cd0
Binary files /dev/null and b/sample_data/sample_docx.docx differ
diff --git a/sample_data/sample_docx.pdf b/sample_data/sample_docx.pdf
new file mode 100644
index 0000000..7890c75
Binary files /dev/null and b/sample_data/sample_docx.pdf differ
diff --git a/sample_data/sample_epub.epub b/sample_data/sample_epub.epub
new file mode 100644
index 0000000..988e401
Binary files /dev/null and b/sample_data/sample_epub.epub differ
diff --git a/sample_data/sample_gif.gif b/sample_data/sample_gif.gif
new file mode 100644
index 0000000..55cd9bf
Binary files /dev/null and b/sample_data/sample_gif.gif differ
diff --git a/sample_data/sample_html.html b/sample_data/sample_html.html
new file mode 100644
index 0000000..80629e6
--- /dev/null
+++ b/sample_data/sample_html.html
@@ -0,0 +1,12 @@
+
+
+
+
+
+ Simple HTML Page
+
+
+ This is a sample HTML file.
+ The answer you're looking for: 32596325
+
+
diff --git a/sample_data/sample_html.pdf b/sample_data/sample_html.pdf
new file mode 100644
index 0000000..a2044dd
Binary files /dev/null and b/sample_data/sample_html.pdf differ
diff --git a/sample_data/sample_jpg.jpg b/sample_data/sample_jpg.jpg
new file mode 100644
index 0000000..4ab3e82
Binary files /dev/null and b/sample_data/sample_jpg.jpg differ
diff --git a/sample_data/sample_json.json b/sample_data/sample_json.json
new file mode 100644
index 0000000..20c6f01
--- /dev/null
+++ b/sample_data/sample_json.json
@@ -0,0 +1,4 @@
+{
+ "Title": "This is a sample JSON file",
+ "Answer": "The answer you're looking for: 485596"
+}
\ No newline at end of file
diff --git a/sample_data/sample_json.pdf b/sample_data/sample_json.pdf
new file mode 100644
index 0000000..25795bc
Binary files /dev/null and b/sample_data/sample_json.pdf differ
diff --git a/sample_data/sample_md.md b/sample_data/sample_md.md
new file mode 100644
index 0000000..aa31c64
--- /dev/null
+++ b/sample_data/sample_md.md
@@ -0,0 +1,2 @@
+**This is a sample .md file.**
+The answer you're looking for: *6565652*
diff --git a/sample_data/sample_md.pdf b/sample_data/sample_md.pdf
new file mode 100644
index 0000000..2ebab67
Binary files /dev/null and b/sample_data/sample_md.pdf differ
diff --git a/sample_data/sample_multi_pdf.pdf b/sample_data/sample_multi_pdf.pdf
new file mode 100644
index 0000000..b462989
Binary files /dev/null and b/sample_data/sample_multi_pdf.pdf differ
diff --git a/sample_data/sample_odp.odp b/sample_data/sample_odp.odp
new file mode 100644
index 0000000..ca8d98c
Binary files /dev/null and b/sample_data/sample_odp.odp differ
diff --git a/sample_data/sample_odp.pdf b/sample_data/sample_odp.pdf
new file mode 100644
index 0000000..5207624
Binary files /dev/null and b/sample_data/sample_odp.pdf differ
diff --git a/sample_data/sample_ods.ods b/sample_data/sample_ods.ods
new file mode 100644
index 0000000..1afda09
Binary files /dev/null and b/sample_data/sample_ods.ods differ
diff --git a/sample_data/sample_ods.pdf b/sample_data/sample_ods.pdf
new file mode 100644
index 0000000..e9b62f6
Binary files /dev/null and b/sample_data/sample_ods.pdf differ
diff --git a/sample_data/sample_odt.odt b/sample_data/sample_odt.odt
new file mode 100644
index 0000000..5e12284
Binary files /dev/null and b/sample_data/sample_odt.odt differ
diff --git a/sample_data/sample_odt.pdf b/sample_data/sample_odt.pdf
new file mode 100644
index 0000000..00eb525
Binary files /dev/null and b/sample_data/sample_odt.pdf differ
diff --git a/sample_data/sample_pdf.pdf b/sample_data/sample_pdf.pdf
new file mode 100644
index 0000000..44e7d4e
Binary files /dev/null and b/sample_data/sample_pdf.pdf differ
diff --git a/sample_data/sample_png.png b/sample_data/sample_png.png
new file mode 100644
index 0000000..deae77f
Binary files /dev/null and b/sample_data/sample_png.png differ
diff --git a/sample_data/sample_ppt.pdf b/sample_data/sample_ppt.pdf
new file mode 100644
index 0000000..a32deb2
Binary files /dev/null and b/sample_data/sample_ppt.pdf differ
diff --git a/sample_data/sample_ppt.ppt b/sample_data/sample_ppt.ppt
new file mode 100644
index 0000000..39dd319
Binary files /dev/null and b/sample_data/sample_ppt.ppt differ
diff --git a/sample_data/sample_pptx.pdf b/sample_data/sample_pptx.pdf
new file mode 100644
index 0000000..c26ac26
Binary files /dev/null and b/sample_data/sample_pptx.pdf differ
diff --git a/sample_data/sample_pptx.pptx b/sample_data/sample_pptx.pptx
new file mode 100644
index 0000000..5132caf
Binary files /dev/null and b/sample_data/sample_pptx.pptx differ
diff --git a/sample_data/sample_rtf.pdf b/sample_data/sample_rtf.pdf
new file mode 100644
index 0000000..9654781
Binary files /dev/null and b/sample_data/sample_rtf.pdf differ
diff --git a/sample_data/sample_rtf.rtf b/sample_data/sample_rtf.rtf
new file mode 100644
index 0000000..4efdb9a
--- /dev/null
+++ b/sample_data/sample_rtf.rtf
@@ -0,0 +1,16 @@
+{\rtf1\ansi\deff3\adeflang1025
+{\fonttbl{\f0\froman\fprq2\fcharset0 Times New Roman;}{\f1\froman\fprq2\fcharset2 Symbol;}{\f2\fswiss\fprq2\fcharset0 Arial;}{\f3\froman\fprq2\fcharset0 Liberation Serif{\*\falt Times New Roman};}{\f4\fswiss\fprq2\fcharset0 Liberation Sans{\*\falt Arial};}{\f5\fnil\fprq2\fcharset0 Noto Sans CJK SC;}{\f6\fswiss\fprq0\fcharset128 Noto Sans Devanagari;}{\f7\fnil\fprq2\fcharset0 Noto Sans Devanagari;}}
+{\colortbl;\red0\green0\blue0;\red0\green0\blue255;\red0\green255\blue255;\red0\green255\blue0;\red255\green0\blue255;\red255\green0\blue0;\red255\green255\blue0;\red255\green255\blue255;\red0\green0\blue128;\red0\green128\blue128;\red0\green128\blue0;\red128\green0\blue128;\red128\green0\blue0;\red128\green128\blue0;\red128\green128\blue128;\red192\green192\blue192;}
+{\stylesheet{\s0\snext0\rtlch\af7\afs24\alang1081 \ltrch\lang1033\langfe2052\hich\af3\loch\widctlpar\hyphpar0\ltrpar\cf0\f3\fs24\lang1033\kerning1\dbch\af8\langfe2052 Normal;}
+{\s15\sbasedon0\snext16\rtlch\af7\afs28 \ltrch\hich\af4\loch\sb240\sa120\keepn\f4\fs28\dbch\af5 Heading;}
+{\s16\sbasedon0\snext16\loch\sl276\slmult1\sb0\sa140 Body Text;}
+{\s17\sbasedon16\snext17\rtlch\af6 \ltrch\loch\sl276\slmult1\sb0\sa140 List;}
+{\s18\sbasedon0\snext18\rtlch\af6\afs24\ai \ltrch\loch\sb120\sa120\noline\fs24\i Caption;}
+{\s19\sbasedon0\snext19\rtlch\af6 \ltrch\loch\noline Index;}
+}{\*\generator LibreOffice/24.2.7.2$Linux_X86_64 LibreOffice_project/420$Build-2}{\info{\creatim\yr2025\mo10\dy24\hr18\min13}{\revtim\yr2025\mo10\dy24\hr18\min14}{\printim\yr0\mo0\dy0\hr0\min0}}{\*\userprops}\deftab709
+\hyphauto1\viewscale180\formshade\nobrkwrptbl\paperh16838\paperw11906\margl1134\margr1134\margt1134\margb1134\sectd\sbknone\sftnnar\saftnnrlc\sectunlocked1\pgwsxn11906\pghsxn16838\marglsxn1134\margrsxn1134\margtsxn1134\margbsxn1134\ftnbj\ftnstart1\ftnrstcont\ftnnar\aenddoc\aftnrstcont\aftnstart1\aftnnrlc
+{\*\ftnsep\chftnsep}\pgndec\pard\plain \s0\rtlch\af7\afs24\alang1081 \ltrch\lang1033\langfe2052\hich\af3\loch\widctlpar\hyphpar0\ltrpar\cf0\f3\fs24\lang1033\kerning1\dbch\af8\langfe2052\ql\ltrpar{\loch
+This is a sample .rtf file.}
+\par \pard\plain \s0\rtlch\af7\afs24\alang1081 \ltrch\lang1033\langfe2052\hich\af3\loch\widctlpar\hyphpar0\ltrpar\cf0\f3\fs24\lang1033\kerning1\dbch\af8\langfe2052\ql\ltrpar{\loch
+The answer you\u8217\'92re looking for: 798945}
+\par }
\ No newline at end of file
diff --git a/sample_data/sample_tif.pdf b/sample_data/sample_tif.pdf
new file mode 100644
index 0000000..8831f42
Binary files /dev/null and b/sample_data/sample_tif.pdf differ
diff --git a/sample_data/sample_tif.tif b/sample_data/sample_tif.tif
new file mode 100644
index 0000000..7ceeab2
Binary files /dev/null and b/sample_data/sample_tif.tif differ
diff --git a/sample_data/sample_txt.pdf b/sample_data/sample_txt.pdf
new file mode 100644
index 0000000..816f721
Binary files /dev/null and b/sample_data/sample_txt.pdf differ
diff --git a/sample_data/sample_txt.txt b/sample_data/sample_txt.txt
new file mode 100644
index 0000000..50fe4d6
--- /dev/null
+++ b/sample_data/sample_txt.txt
@@ -0,0 +1,2 @@
+This is a sample .txt file.
+The answer you're looking for: 2676599
\ No newline at end of file
diff --git a/sample_data/sample_xls.pdf b/sample_data/sample_xls.pdf
new file mode 100644
index 0000000..8f92d8d
Binary files /dev/null and b/sample_data/sample_xls.pdf differ
diff --git a/sample_data/sample_xls.xls b/sample_data/sample_xls.xls
new file mode 100644
index 0000000..399cf2a
Binary files /dev/null and b/sample_data/sample_xls.xls differ
diff --git a/sample_data/sample_xlsx.pdf b/sample_data/sample_xlsx.pdf
new file mode 100644
index 0000000..51f8320
Binary files /dev/null and b/sample_data/sample_xlsx.pdf differ
diff --git a/sample_data/sample_xlsx.xlsx b/sample_data/sample_xlsx.xlsx
new file mode 100644
index 0000000..a79bb5b
Binary files /dev/null and b/sample_data/sample_xlsx.xlsx differ
diff --git a/sample_data/sample_yaml.pdf b/sample_data/sample_yaml.pdf
new file mode 100644
index 0000000..0e98b51
Binary files /dev/null and b/sample_data/sample_yaml.pdf differ
diff --git a/sample_data/sample_yaml.yaml b/sample_data/sample_yaml.yaml
new file mode 100644
index 0000000..5f167d1
--- /dev/null
+++ b/sample_data/sample_yaml.yaml
@@ -0,0 +1,2 @@
+Title: "This is a sample YAML file"
+Answer: "The answer you're looking for: 559585"
\ No newline at end of file
diff --git a/scripts/benchmark.py b/scripts/benchmark.py
new file mode 100644
index 0000000..8887cd0
--- /dev/null
+++ b/scripts/benchmark.py
@@ -0,0 +1,367 @@
+"""
+Benchmark script for comparing embedding server solutions.
+
+Usage (Sol A — vLLM):
+ uv run python scripts/benchmark.py \
+ --solution a \
+ --server-url http://localhost:18000 \
+ --model athrael-soju/colqwen3.5-4.5B-v3 \
+ --data-dir ../toy_data/smartcockpit \
+ --output benchmark_results/sol_a.json
+
+Usage (Sol B — custom FastAPI):
+ uv run python scripts/benchmark.py \
+ --solution b \
+ --server-url http://localhost:18001 \
+ --model vidore/colqwen2-v1.0 \
+ --data-dir ../toy_data/smartcockpit \
+ --output benchmark_results/sol_b.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional
+
+# ---------------------------------------------------------------------------
+# Queries — precise aviation questions whose answers are sentence/paragraph
+# form, drawn from A319/A320 content in the smartcockpit PDFs.
+# ---------------------------------------------------------------------------
+QUERIES = [
+ "What is the normal operating cabin altitude in cruise for the A320?",
+ "How is the bleed air supply to the air conditioning system controlled on the A320?",
+ "What happens to pressurization if both outflow valves fail in the open position?",
+ "What is the purpose of the Ram Air inlet on the A320 air conditioning system?",
+ "Describe the function of the Pack Flow Control Valve and its operating modes.",
+]
+
+QUERY_RUNS = 5 # repeat each query N times to get stable latency
+
+
+# ---------------------------------------------------------------------------
+# GPU VRAM sampler (non-blocking, best-effort)
+# ---------------------------------------------------------------------------
+
+def _sample_vram_mb() -> Optional[float]:
+ """Return total used VRAM in MB across all GPUs, or None if unavailable."""
+ try:
+ out = subprocess.check_output(
+ ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
+ stderr=subprocess.DEVNULL,
+ timeout=5,
+ )
+ values = [int(x.strip()) for x in out.decode().splitlines() if x.strip()]
+ return float(sum(values)) if values else None
+ except Exception:
+ return None
+
+
+# ---------------------------------------------------------------------------
+# Sol B client (httpx-based, mirrors feature/client-server transport)
+# ---------------------------------------------------------------------------
+
+def _sol_b_health(server_url: str) -> bool:
+ import httpx
+ try:
+ r = httpx.get(f"{server_url.rstrip('/')}/health", timeout=5)
+ return r.status_code == 200
+ except Exception:
+ return False
+
+
+def _sol_b_embed_images(server_url: str, images, model_name: str):
+ """Call Sol B POST /v1/embed/images using torch binary transport."""
+ import io
+ import torch
+ import httpx
+
+ buf = io.BytesIO()
+ # images: list of PIL.Image → convert to PNG bytes
+ image_bytes_list = []
+ for img in images:
+ ibuf = io.BytesIO()
+ img.save(ibuf, format="PNG")
+ image_bytes_list.append(ibuf.getvalue())
+
+ payload = {"model": model_name, "images": image_bytes_list}
+ torch.save(payload, buf)
+ buf.seek(0)
+
+ with httpx.Client(timeout=600) as client:
+ resp = client.post(
+ f"{server_url.rstrip('/')}/v1/embed/images",
+ content=buf.read(),
+ headers={"Content-Type": "application/octet-stream"},
+ )
+ resp.raise_for_status()
+
+ result_buf = io.BytesIO(resp.content)
+ result = torch.load(result_buf, weights_only=False)
+ return result["embeddings"]
+
+
+def _sol_b_embed_query(server_url: str, query: str, model_name: str):
+ """Call Sol B POST /v1/embed/queries using torch binary transport."""
+ import io
+ import torch
+ import httpx
+
+ payload = {"model": model_name, "queries": [query]}
+ buf = io.BytesIO()
+ torch.save(payload, buf)
+ buf.seek(0)
+
+ with httpx.Client(timeout=600) as client:
+ resp = client.post(
+ f"{server_url.rstrip('/')}/v1/embed/queries",
+ content=buf.read(),
+ headers={"Content-Type": "application/octet-stream"},
+ )
+ resp.raise_for_status()
+
+ result_buf = io.BytesIO(resp.content)
+ result = torch.load(result_buf, weights_only=False)
+ return result["embeddings"]
+
+
+# ---------------------------------------------------------------------------
+# Main benchmark logic
+# ---------------------------------------------------------------------------
+
+def run_sol_a(server_url: str, model_name: str, data_dir: Path, index_root: Path) -> dict:
+ """Benchmark Solution A: vLLM-backed MultiModalRetrieverModel."""
+ from foretrieval import MultiModalRetrieverModel
+ from foretrieval.embedding_server import EmbeddingServerConfig
+
+ cfg = EmbeddingServerConfig(
+ url=server_url,
+ model_name=model_name,
+ auto_deploy=False,
+ )
+
+ # Health check
+ from foretrieval.embedding_server.client import EmbeddingServerClient
+ client = EmbeddingServerClient(cfg)
+ if not client.health_check():
+ print(f"ERROR: Sol A server not healthy at {server_url}", file=sys.stderr)
+ sys.exit(1)
+ print(f"Sol A server healthy at {server_url}")
+
+ # --- Indexing ---
+ print("Indexing smartcockpit corpus (Sol A)...")
+ pdf_files = list(data_dir.glob("*.pdf")) + list(data_dir.glob("*.PDF"))
+ n_docs = len(pdf_files)
+ vram_before = _sample_vram_mb()
+ t0 = time.perf_counter()
+ model = MultiModalRetrieverModel.from_pretrained(
+ model_name,
+ index_root=str(index_root),
+ device="cpu", # no local GPU needed — embeddings go to server
+ verbose=1,
+ embedding_server=cfg,
+ )
+ model.index(str(data_dir), index_name="smartcockpit_bench", overwrite=True)
+ index_time = time.perf_counter() - t0
+ vram_after_index = _sample_vram_mb()
+ print(f" Indexing done in {index_time:.1f}s")
+
+ # --- Query latency ---
+ query_results = []
+ for q in QUERIES:
+ latencies = []
+ top_pages = None
+ for run in range(QUERY_RUNS):
+ t0 = time.perf_counter()
+ results = model.search(q, k=3)
+ latencies.append(time.perf_counter() - t0)
+ if run == 0:
+ top_pages = [
+ {"doc_id": r.doc_id, "page_num": r.page_num, "score": float(r.score) if r.score is not None else None}
+ for r in results
+ ] if results else []
+ mean_lat = sum(latencies) / len(latencies)
+ std_lat = (sum((x - mean_lat) ** 2 for x in latencies) / len(latencies)) ** 0.5
+ print(f" Query '{q[:60]}...' → mean {mean_lat:.3f}s ± {std_lat:.3f}s")
+ query_results.append({
+ "query": q,
+ "latency_mean_s": round(mean_lat, 4),
+ "latency_std_s": round(std_lat, 4),
+ "latency_all_s": [round(x, 4) for x in latencies],
+ "top_3_pages": top_pages,
+ })
+
+ vram_peak = _sample_vram_mb()
+
+ return {
+ "solution": "A",
+ "server_url": server_url,
+ "model_name": model_name,
+ "n_docs": n_docs,
+ "indexing_time_s": round(index_time, 2),
+ "vram_before_index_mb": vram_before,
+ "vram_after_index_mb": vram_after_index,
+ "vram_peak_mb": vram_peak,
+ "queries": query_results,
+ }
+
+
+def run_sol_b(server_url: str, model_name: str, data_dir: Path) -> dict:
+ """
+ Benchmark Solution B: custom FastAPI server.
+
+ Sol B has no MultiModalRetrieverModel integration so we call the HTTP
+ endpoints directly and measure raw embedding throughput. We convert
+ PDFs to page images (same pipeline FORetrieval uses) then embed them.
+ """
+ try:
+ import httpx # noqa: F401
+ import torch # noqa: F401
+ except ImportError as e:
+ print(f"ERROR: missing dep for Sol B benchmark: {e}", file=sys.stderr)
+ sys.exit(1)
+
+ # Health check
+ if not _sol_b_health(server_url):
+ print(f"ERROR: Sol B server not healthy at {server_url}", file=sys.stderr)
+ sys.exit(1)
+ print(f"Sol B server healthy at {server_url}")
+
+ # Warm-up: trigger model load before timing (Sol B loads lazily on first request)
+ print("Warming up Sol B server (model load)...")
+ from PIL import Image as PILImage
+ _warmup_img = PILImage.new("RGB", (64, 64), color=(128, 128, 128))
+ t_warmup = time.perf_counter()
+ _sol_b_embed_images(server_url, [_warmup_img], model_name)
+ print(f" Warm-up done in {time.perf_counter() - t_warmup:.1f}s (model load included)")
+
+ # Convert PDFs to images (reuse pdf2image, same as FORetrieval internals)
+ from pdf2image import convert_from_path
+
+ pdf_files = sorted((list(data_dir.glob("*.pdf")) + list(data_dir.glob("*.PDF"))))
+ print(f"Converting {len(pdf_files)} PDFs to images...")
+ all_images = []
+ for pdf in pdf_files:
+ pages = convert_from_path(str(pdf), dpi=150)
+ all_images.extend(pages)
+ print(f" {pdf.name}: {len(pages)} pages")
+ n_pages = len(all_images)
+ n_docs = len(pdf_files)
+ print(f"Total pages: {n_pages}")
+
+ # --- Indexing (embedding all pages) ---
+ print("Embedding all pages (Sol B)...")
+ vram_before = _sample_vram_mb()
+ BATCH = 4
+ t0 = time.perf_counter()
+ all_embeddings = []
+ for i in range(0, n_pages, BATCH):
+ batch = all_images[i : i + BATCH]
+ emb = _sol_b_embed_images(server_url, batch, model_name)
+ all_embeddings.append(emb)
+ print(f" Embedded pages {i+1}-{min(i+BATCH, n_pages)}/{n_pages}")
+ index_time = time.perf_counter() - t0
+ vram_after_index = _sample_vram_mb()
+ print(f" Indexing done in {index_time:.1f}s ({n_pages / index_time:.2f} pages/s)")
+
+ # --- Query latency ---
+ query_results = []
+ for q in QUERIES:
+ latencies = []
+ for run in range(QUERY_RUNS):
+ t0 = time.perf_counter()
+ _sol_b_embed_query(server_url, q, model_name)
+ latencies.append(time.perf_counter() - t0)
+ mean_lat = sum(latencies) / len(latencies)
+ std_lat = (sum((x - mean_lat) ** 2 for x in latencies) / len(latencies)) ** 0.5
+ print(f" Query '{q[:60]}...' → mean {mean_lat:.3f}s ± {std_lat:.3f}s")
+ query_results.append({
+ "query": q,
+ "latency_mean_s": round(mean_lat, 4),
+ "latency_std_s": round(std_lat, 4),
+ "latency_all_s": [round(x, 4) for x in latencies],
+ "top_3_pages": None, # no retrieval integration in Sol B direct mode
+ })
+
+ vram_peak = _sample_vram_mb()
+
+ return {
+ "solution": "B",
+ "server_url": server_url,
+ "model_name": model_name,
+ "n_docs": n_docs,
+ "n_pages": n_pages,
+ "indexing_time_s": round(index_time, 2),
+ "pages_per_sec": round(n_pages / index_time, 2),
+ "vram_before_index_mb": vram_before,
+ "vram_after_index_mb": vram_after_index,
+ "vram_peak_mb": vram_peak,
+ "queries": query_results,
+ }
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+def main():
+ parser = argparse.ArgumentParser(description="Embedding server benchmark")
+ parser.add_argument("--solution", choices=["a", "b"], required=True,
+ help="Which solution to benchmark: 'a' (vLLM) or 'b' (custom FastAPI)")
+ parser.add_argument("--server-url", required=True,
+ help="Base URL of the embedding server (e.g. http://localhost:18000)")
+ parser.add_argument("--model", required=True,
+ help="Model name served by the server")
+ parser.add_argument("--data-dir", default="../toy_data/smartcockpit",
+ help="Path to directory containing smartcockpit PDFs")
+ parser.add_argument("--output", default=None,
+ help="Path to write JSON results (default: benchmark_results/sol_{a,b}.json)")
+ args = parser.parse_args()
+
+ data_dir = Path(args.data_dir).resolve()
+ if not data_dir.exists():
+ print(f"ERROR: data-dir not found: {data_dir}", file=sys.stderr)
+ sys.exit(1)
+
+ out_path = Path(args.output) if args.output else Path(f"benchmark_results/sol_{args.solution}.json")
+ out_path.parent.mkdir(parents=True, exist_ok=True)
+
+ print(f"=== Benchmark Solution {'A (vLLM)' if args.solution == 'a' else 'B (custom FastAPI)'} ===")
+ print(f"Server : {args.server_url}")
+ print(f"Model : {args.model}")
+ print(f"Data : {data_dir}")
+ print()
+
+ if args.solution == "a":
+ index_root = Path(tempfile.mkdtemp(prefix="foretrieval_bench_"))
+ try:
+ results = run_sol_a(args.server_url, args.model, data_dir, index_root)
+ finally:
+ shutil.rmtree(index_root, ignore_errors=True)
+ else:
+ results = run_sol_b(args.server_url, args.model, data_dir)
+
+ out_path.write_text(json.dumps(results, indent=2))
+ print(f"\nResults written to {out_path}")
+
+ # Print summary
+ print("\n--- Summary ---")
+ print(f"Indexing time : {results['indexing_time_s']}s")
+ if "pages_per_sec" in results:
+ print(f"Pages/sec : {results['pages_per_sec']}")
+ if results.get("vram_peak_mb"):
+ print(f"Peak VRAM : {results['vram_peak_mb']} MB")
+ q_means = [q["latency_mean_s"] for q in results["queries"]]
+ overall_mean = sum(q_means) / len(q_means)
+ print(f"Query latency : {overall_mean:.3f}s avg over {len(QUERIES)} queries x {QUERY_RUNS} runs")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/all.py b/tests/all.py
deleted file mode 100644
index 85ad522..0000000
--- a/tests/all.py
+++ /dev/null
@@ -1,170 +0,0 @@
-from pathlib import Path
-
-from colpali_engine.utils.torch_utils import get_torch_device
-
-from byaldi import RAGMultiModalModel
-
-device = get_torch_device("auto")
-print(f"Using device: {device}")
-
-path_document_1 = Path("docs/attention.pdf")
-path_document_2 = Path("docs/attention_copy.pdf")
-
-
-def test_single_pdf():
- print("Testing single PDF indexing and retrieval...")
-
- # Initialize the model
- model = RAGMultiModalModel.from_pretrained("vidore/colpali-v1.2", device=device)
-
- if not Path("docs/attention.pdf").is_file():
- raise FileNotFoundError(
- f"Please download the PDF file from https://arxiv.org/pdf/1706.03762 and move it to {path_document_1}."
- )
-
- # Index a single PDF
- model.index(
- input_path="docs/attention.pdf",
- index_name="attention_index",
- store_collection_with_index=True,
- overwrite=True,
- )
-
- # Test retrieval
- queries = [
- "How does the positional encoding thing work?",
- "what's the BLEU score of this new strange method?",
- ]
-
- for query in queries:
- results = model.search(query, k=3)
-
- print(f"\nQuery: {query}")
- for result in results:
- print(
- f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
- )
-
- # Check if the expected page (6 for positional encoding) is in the top results
- if "positional encoding" in query.lower():
- assert any(
- r.page_num == 6 for r in results
- ), "Expected page 6 for positional encoding query"
-
- # Check if the expected pages (8 and 9 for BLEU score) are in the top results
- if "bleu score" in query.lower():
- assert any(
- r.page_num in [8, 9] for r in results
- ), "Expected pages 8 or 9 for BLEU score query"
-
- print("Single PDF test completed.")
-
-
-def test_multi_document():
- print("\nTesting multi-document indexing and retrieval...")
-
- # Initialize the model
- model = RAGMultiModalModel.from_pretrained("vidore/colpali")
-
- if not Path("docs/attention.pdf").is_file():
- raise FileNotFoundError(
- f"Please download the PDF file from https://arxiv.org/pdf/1706.03762 and move it to {path_document_1}."
- )
- if not Path("docs/attention_copy.pdf").is_file():
- raise FileNotFoundError(
- f"Please download the PDF file from https://arxiv.org/pdf/1706.03762 and move it to {path_document_2}."
- )
-
- # Index a directory of documents
- model.index(
- input_path="docs/",
- index_name="multi_doc_index",
- store_collection_with_index=True,
- overwrite=True,
- )
-
- # Test retrieval
- queries = [
- "How does the positional encoding thing work?",
- "what's the BLEU score of this new strange method?",
- ]
-
- for query in queries:
- results = model.search(query, k=5)
-
- print(f"\nQuery: {query}")
- for result in results:
- print(
- f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
- )
-
- # Check if the expected page (6 for positional encoding) is in the top results
- if "positional encoding" in query.lower():
- assert any(
- r.page_num == 6 for r in results
- ), "Expected page 6 for positional encoding query"
-
- # Check if the expected pages (8 and 9 for BLEU score) are in the top results
- if "bleu score" in query.lower():
- assert any(
- r.page_num in [8, 9] for r in results
- ), "Expected pages 8 or 9 for BLEU score query"
-
- print("Multi-document test completed.")
-
-
-def test_add_to_index():
- print("\nTesting adding to an existing index...")
-
- # Load the existing index
- model = RAGMultiModalModel.from_index("multi_doc_index")
-
- # Add a new document to the index
- model.add_to_index(
- input_item="docs/",
- store_collection_with_index=True,
- doc_id=[1002, 1003],
- metadata=[{"author": "John Doe", "year": 2023}] * 2,
- )
-
- # Test retrieval with the updated index
- queries = ["what's the BLEU score of this new strange method?"]
-
- for query in queries:
- results = model.search(query, k=3)
-
- print(f"\nQuery: {query}")
- for result in results:
- print(
- f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
- )
- print(f"Metadata: {result.metadata}")
-
- # Check if the expected page (6 for positional encoding) is in the top results
- if "positional encoding" in query.lower():
- assert any(
- r.page_num == 6 for r in results
- ), "Expected page 6 for positional encoding query"
-
- # Check if the expected pages (8 and 9 for BLEU score) are in the top results
- if "bleu score" in query.lower():
- assert any(
- r.page_num in [8, 9] for r in results
- ), "Expected pages 8 or 9 for BLEU score query"
-
- print("Add to index test completed.")
-
-
-if __name__ == "__main__":
- print("Starting tests...")
-
- print("/n/n----------------- Single PDF test -----------------n")
- test_single_pdf()
-
- print("/n/n----------------- Multi document test -----------------n")
- test_multi_document()
-
- print("/n/n----------------- Add to index test -----------------n")
- test_add_to_index()
-
- print("\nAll tests completed.")
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..8df8087
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,93 @@
+"""
+Shared fixtures for the FORetrieval test suite.
+"""
+
+from __future__ import annotations
+
+import os
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+
+TESTS_DIR = Path(__file__).parent
+DATA_DIR = TESTS_DIR / "data"
+SAMPLE_DATA_DIR = Path(__file__).parent.parent / "sample_data"
+
+# ---------------------------------------------------------------------------
+# AI backend selection — same priority order as FORag
+# ---------------------------------------------------------------------------
+
+_BACKEND_MAP = {
+ "openrouter": ("OPENROUTER_API_KEY", "mistralai/mistral-small-3.2-24b-instruct"),
+ "openai": ("OPENAI_API_KEY", "gpt-4o-mini"),
+ "mistral": ("MISTRAL_API_KEY", "mistral-small-latest"),
+}
+
+
+def _all_available_ai_backends() -> list[dict]:
+ """Return an ai_cfg dict for every available backend.
+
+ For cloud backends the credential is an API key. For Ollama, the
+ ``base_url`` key points to the daemon's OpenAI-compatible endpoint.
+
+ Priority order: OpenRouter → OpenAI → Mistral → Ollama.
+ This mirrors the backend selection used in FORag's integration suite.
+
+ Ollama model defaults to ``mistral-small-latest`` (a text-only model
+ suitable for metadata generation). Override with the ``OLLAMA_MODEL``
+ environment variable.
+ """
+ available = []
+
+ for provider, (env_var, model) in _BACKEND_MAP.items():
+ key = os.environ.get(env_var)
+ if key:
+ available.append({"provider": provider, "name": model, "api_key": key})
+
+ host = os.environ.get("OLLAMA_HOST")
+ if host:
+ try:
+ urllib.request.urlopen(f"{host}/api/tags", timeout=3)
+ # Default to a text-only model — not a vision model — since
+ # metadata generation uses text prompts only.
+ model = os.environ.get("OLLAMA_MODEL", "mistral-small-latest")
+ available.append(
+ {"provider": "ollama", "name": model, "base_url": f"{host}/v1"}
+ )
+ except (urllib.error.URLError, OSError):
+ pass # daemon unreachable — skip silently
+
+ return available
+
+
+@pytest.fixture(
+ scope="module",
+ params=_all_available_ai_backends(),
+ ids=lambda cfg: cfg["provider"],
+)
+def ai_cfg(request) -> dict:
+ """Module-scoped fixture parametrized over every available AI backend.
+
+ Each test module that depends on this fixture will be run once per
+ available backend. If no backend is available the entire module is
+ skipped.
+
+ Set one of the following environment variables to enable AI tests:
+
+ OPENROUTER_API_KEY (preferred — fast, high rate limits)
+ OPENAI_API_KEY
+ MISTRAL_API_KEY
+ OLLAMA_HOST (+ optionally OLLAMA_MODEL, default: mistral-small-latest)
+ """
+ if not _all_available_ai_backends():
+ pytest.skip(
+ "No API key or Ollama daemon found. Set OPENROUTER_API_KEY, "
+ "OPENAI_API_KEY, MISTRAL_API_KEY, or OLLAMA_HOST."
+ )
+ return request.param
diff --git a/tests/data/lm317_voltage_regulator.pdf b/tests/data/lm317_voltage_regulator.pdf
new file mode 100644
index 0000000..ed88174
Binary files /dev/null and b/tests/data/lm317_voltage_regulator.pdf differ
diff --git a/tests/data/uln2003_driver.pdf b/tests/data/uln2003_driver.pdf
new file mode 100644
index 0000000..2673fc8
Binary files /dev/null and b/tests/data/uln2003_driver.pdf differ
diff --git a/tests/smoke_test_backends.py b/tests/smoke_test_backends.py
new file mode 100644
index 0000000..2eb1888
--- /dev/null
+++ b/tests/smoke_test_backends.py
@@ -0,0 +1,80 @@
+"""
+Smoke test: index toy_data/smartcockpit with all three backends using a remote embedding server,
+then run one query and compare top-1 results.
+
+Usage:
+ uv run python FORetrieval/tests/smoke_test_backends.py
+"""
+from __future__ import annotations
+
+import os
+import sys
+import tempfile
+from pathlib import Path
+
+# Add FORetrieval to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent / "FORetrieval"))
+
+from foretrieval import MultiModalRetrieverModel
+from foretrieval.embedding_server import EmbeddingServerConfig
+
+EMBEDDING_SERVER_URL = os.getenv("EMBEDDING_SERVER_URL", "http://localhost:18000")
+MODEL_NAME = "athrael-soju/colqwen3.5-4.5B-v3"
+DATA_DIR = Path(__file__).parent.parent.parent / "toy_data" / "smartcockpit"
+QUERY = "What is the normal operating cabin altitude in cruise for the A320?"
+TOP_K = 3
+
+embedding_cfg = EmbeddingServerConfig(
+ url=EMBEDDING_SERVER_URL,
+ model_name=MODEL_NAME,
+ auto_deploy=False,
+ batch_size=4,
+)
+
+BACKENDS = ["local", "qdrant", "milvus"]
+results_by_backend = {}
+
+with tempfile.TemporaryDirectory() as tmp:
+ for backend in BACKENDS:
+ print(f"\n{'='*60}")
+ print(f"Backend: {backend}")
+ print(f"{'='*60}")
+
+ rag = MultiModalRetrieverModel.from_pretrained(
+ pretrained_model_name_or_path=MODEL_NAME,
+ index_root=tmp,
+ storage_backend=backend,
+ embedding_server=embedding_cfg,
+ device="cuda",
+ verbose=1,
+ )
+
+ rag.index(
+ input_path=str(DATA_DIR),
+ index_name=f"smoke_{backend}",
+ overwrite=True,
+ )
+
+ print(f"Query: {QUERY}")
+ results = rag.search(QUERY, k=TOP_K, return_base64_results=False)
+ results_by_backend[backend] = results
+
+ for i, r in enumerate(results):
+ print(f" [{i+1}] doc_id={r.doc_id} page={r.page_num} score={r.score:.4f}")
+
+# Compare top-1 file across backends
+print("\n" + "="*60)
+print("Comparison summary")
+print("="*60)
+print(f"Query: {QUERY}\n")
+for backend, results in results_by_backend.items():
+ if results:
+ r = results[0]
+ print(f" {backend:6s}: doc_id={r.doc_id} page={r.page_num} score={r.score:.4f}")
+ else:
+ print(f" {backend:6s}: NO RESULTS")
+
+# All backends should return at least 1 result
+for backend, results in results_by_backend.items():
+ assert len(results) > 0, f"{backend} returned no results!"
+print("\nAll backends returned results. Smoke test PASSED.")
diff --git a/tests/test_colpali.py b/tests/test_colpali.py
deleted file mode 100644
index f8e85ae..0000000
--- a/tests/test_colpali.py
+++ /dev/null
@@ -1,23 +0,0 @@
-from typing import Generator
-
-import pytest
-from colpali_engine.models import ColPali
-from colpali_engine.utils.torch_utils import get_torch_device, tear_down_torch
-
-from byaldi import RAGMultiModalModel
-from byaldi.colpali import ColPaliModel
-
-
-@pytest.fixture(scope="module")
-def colpali_rag_model() -> Generator[RAGMultiModalModel, None, None]:
- device = get_torch_device("auto")
- print(f"Using device: {device}")
- yield RAGMultiModalModel.from_pretrained("vidore/colpali-v1.2", device=device)
- tear_down_torch()
-
-
-@pytest.mark.slow
-def test_load_colpali_from_pretrained(colpali_rag_model: RAGMultiModalModel):
- assert isinstance(colpali_rag_model, RAGMultiModalModel)
- assert isinstance(colpali_rag_model.model, ColPaliModel)
- assert isinstance(colpali_rag_model.model.model, ColPali)
diff --git a/tests/test_colpali_backend_dispatch.py b/tests/test_colpali_backend_dispatch.py
new file mode 100644
index 0000000..2039f31
--- /dev/null
+++ b/tests/test_colpali_backend_dispatch.py
@@ -0,0 +1,407 @@
+"""Tests for ColPaliModel backend dispatch via the VectorStore interface.
+
+These tests verify that:
+- storage_backend="local/qdrant/milvus" selects the correct VectorStore class
+- the deprecated storage_qdrant boolean still works with a DeprecationWarning
+- backward-compat properties (storage_qdrant, qdrant_client, qdrant_collection,
+ indexed_embeddings, embed_id_to_doc_id) still work
+
+No GPU or real model is required — ColPaliModel.__init__ is bypassed via
+direct attribute injection on MagicMock objects.
+"""
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from foretrieval.vector_store import (
+ LocalVectorStore,
+ QdrantVectorStore,
+ MilvusVectorStore,
+ make_vector_store,
+)
+from foretrieval.vector_store.base import make_point_id, StoredPoint
+
+
+# ---------------------------------------------------------------------------
+# Helpers — build a ColPaliModel with all heavy parts mocked
+# ---------------------------------------------------------------------------
+
+def _make_model(storage_backend: str = "local", **extra_kwargs):
+ """Build a ColPaliModel instance with GPU/model loading patched out."""
+ with (
+ patch("foretrieval.colpali.ColPaliModel._load_model_and_processor"),
+ patch("foretrieval.colpali.ColPaliModel._load_processor_only"),
+ ):
+ from foretrieval.colpali import ColPaliModel
+
+ model = ColPaliModel.__new__(ColPaliModel)
+
+ # Minimal attributes to avoid AttributeError from __init__ side-effects
+ model.pretrained_model_name_or_path = "vidore/colpali-v1.2-test"
+ model.model_name = "vidore/colpali-v1.2-test"
+ model.verbose = 0
+ model.load_from_index = False
+ model.index_root = ".foretrieval_test"
+ model.index_name = None
+ model.kwargs = {}
+ model.storage_backend = storage_backend.strip().lower()
+ model.storage_config = extra_kwargs.get("storage_config", {})
+ model._storage_qdrant_compat = (model.storage_backend == "qdrant")
+ model.ingestion = {"backend": "default"}
+ model.ingestion_backend = "default"
+ model.n_gpu = 0
+ model.device = "cpu"
+ model.load_in_4bit = False
+ model.load_in_8bit = False
+ model.bnb_4bit_quant_type = "nf4"
+ model.bnb_4bit_compute_dtype = "float16"
+ model.collection = {}
+ model.embed_id_to_extra = {}
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.doc_ids = set()
+ model.enable_heatmaps = False
+ model.enable_circle = False
+ model.full_document_collection = False
+ model.resize_stored_images = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.docling_dir = None
+ model.SOURCE_EXTS = set()
+ model.IMAGE_EXTS = set()
+ model._remote_client = None
+ model.model = None
+ model.processor = MagicMock()
+
+ # Build real VectorStore
+ model.vector_store = make_vector_store(storage_backend, model.storage_config)
+
+ return model
+
+
+# ---------------------------------------------------------------------------
+# Backend dispatch
+# ---------------------------------------------------------------------------
+
+class TestBackendDispatch:
+ def test_local_backend(self):
+ model = _make_model("local")
+ assert isinstance(model.vector_store, LocalVectorStore)
+
+ def test_qdrant_backend(self):
+ model = _make_model("qdrant")
+ assert isinstance(model.vector_store, QdrantVectorStore)
+
+ def test_milvus_backend(self):
+ model = _make_model("milvus")
+ assert isinstance(model.vector_store, MilvusVectorStore)
+
+ def test_milvus_candidate_limit_forwarded(self):
+ model = _make_model("milvus", storage_config={"candidate_limit": 128})
+ assert isinstance(model.vector_store, MilvusVectorStore)
+ assert model.vector_store._candidate_limit == 128
+
+
+# ---------------------------------------------------------------------------
+# Deprecated storage_qdrant shim
+# ---------------------------------------------------------------------------
+
+class TestDeprecatedStorageQdrant:
+ def test_storage_qdrant_true_maps_to_qdrant_backend(self):
+ """ColPaliModel(storage_qdrant=True) should emit DeprecationWarning and select qdrant."""
+ with (
+ patch("foretrieval.colpali.ColPaliModel._load_model_and_processor"),
+ patch("foretrieval.colpali.ColPaliModel._load_processor_only"),
+ ):
+ from foretrieval.colpali import ColPaliModel
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ model = ColPaliModel(
+ pretrained_model_name_or_path="vidore/colpali-v1.2",
+ storage_qdrant=True,
+ device="cpu",
+ verbose=0,
+ )
+ # Should have emitted exactly one DeprecationWarning
+ dep_warnings = [x for x in w if issubclass(x.category, DeprecationWarning)]
+ assert len(dep_warnings) == 1
+ assert "storage_qdrant" in str(dep_warnings[0].message)
+ assert model.storage_backend == "qdrant"
+
+ def test_storage_qdrant_false_maps_to_local_backend(self):
+ with (
+ patch("foretrieval.colpali.ColPaliModel._load_model_and_processor"),
+ patch("foretrieval.colpali.ColPaliModel._load_processor_only"),
+ ):
+ from foretrieval.colpali import ColPaliModel
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ model = ColPaliModel(
+ pretrained_model_name_or_path="vidore/colpali-v1.2",
+ storage_qdrant=False,
+ device="cpu",
+ verbose=0,
+ )
+ assert model.storage_backend == "local"
+
+
+# ---------------------------------------------------------------------------
+# Backward-compat property accessors
+# ---------------------------------------------------------------------------
+
+class TestBackwardCompatProperties:
+ def test_storage_qdrant_property_true_for_qdrant_backend(self):
+ model = _make_model("qdrant")
+ assert model.storage_qdrant is True
+
+ def test_storage_qdrant_property_false_for_local_backend(self):
+ model = _make_model("local")
+ assert model.storage_qdrant is False
+
+ def test_storage_qdrant_property_false_for_milvus_backend(self):
+ model = _make_model("milvus")
+ assert model.storage_qdrant is False
+
+ def test_qdrant_client_property_returns_none_for_local(self):
+ model = _make_model("local")
+ assert model.qdrant_client is None
+
+ def test_qdrant_client_property_returns_none_for_milvus(self):
+ model = _make_model("milvus")
+ assert model.qdrant_client is None
+
+ def test_indexed_embeddings_local_backend_returns_list(self):
+ model = _make_model("local")
+ assert isinstance(model.indexed_embeddings, list)
+
+ def test_indexed_embeddings_qdrant_backend_returns_empty(self):
+ model = _make_model("qdrant")
+ assert model.indexed_embeddings == []
+
+ def test_embed_id_to_doc_id_local_backend_returns_dict(self):
+ model = _make_model("local")
+ assert isinstance(model.embed_id_to_doc_id, dict)
+
+ def test_embed_id_to_doc_id_qdrant_backend_returns_empty(self):
+ model = _make_model("qdrant")
+ assert model.embed_id_to_doc_id == {}
+
+
+# ---------------------------------------------------------------------------
+# VectorStore open() called with correct args
+# ---------------------------------------------------------------------------
+
+class TestVectorStoreOpenCall:
+ def test_open_called_with_index_name(self, tmp_path):
+ """When ColPaliModel.index() is called, vector_store.open() gets the index name."""
+ model = _make_model("local")
+ model.index_root = str(tmp_path)
+
+ # Patch open on the LocalVectorStore so we can verify the call
+ from foretrieval.vector_store.local import LocalVectorStore as LSV
+ original_open = LSV.open
+ calls = []
+
+ def fake_open(self_, index_name, index_root, *, create, dim=None):
+ calls.append((index_name, create))
+ original_open(self_, index_name, index_root, create=create, dim=dim)
+
+ with patch.object(LSV, "open", fake_open):
+ model.vector_store = make_vector_store("local")
+ model.vector_store.open("my_index", tmp_path, create=True)
+
+ assert ("my_index", True) in calls
+
+
+# ---------------------------------------------------------------------------
+# index_config storage_backend field
+# ---------------------------------------------------------------------------
+
+class TestIndexConfigStorageBackend:
+ """Verify _export_index writes the correct storage_backend string."""
+
+ def test_export_writes_local(self, tmp_path):
+ import srsly
+ model = _make_model("local")
+ model.index_root = str(tmp_path)
+ model.index_name = "test_export"
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.embed_id_to_extra = {}
+ model.full_document_collection = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.model_name = "vidore/colpali-v1.2"
+
+ model._export_index()
+
+ cfg = srsly.read_gzip_json(tmp_path / "test_export" / "index_config.json.gz")
+ assert cfg["storage_backend"] == "local"
+
+ def test_export_writes_qdrant(self, tmp_path):
+ import srsly
+ model = _make_model("qdrant")
+ model.index_root = str(tmp_path)
+ model.index_name = "test_export_qdrant"
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.embed_id_to_extra = {}
+ model.full_document_collection = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.model_name = "vidore/colpali-v1.2"
+
+ # Patch vector_store.export_sidecar to avoid Qdrant client calls
+ model.vector_store = MagicMock()
+ model._export_index()
+
+ cfg = srsly.read_gzip_json(tmp_path / "test_export_qdrant" / "index_config.json.gz")
+ assert cfg["storage_backend"] == "qdrant"
+
+ def test_export_writes_milvus(self, tmp_path):
+ import srsly
+ model = _make_model("milvus")
+ model.index_root = str(tmp_path)
+ model.index_name = "test_export_milvus"
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.embed_id_to_extra = {}
+ model.full_document_collection = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.model_name = "vidore/colpali-v1.2"
+
+ model.vector_store = MagicMock()
+ model._export_index()
+
+ cfg = srsly.read_gzip_json(tmp_path / "test_export_milvus" / "index_config.json.gz")
+ assert cfg["storage_backend"] == "milvus"
+
+
+# ---------------------------------------------------------------------------
+# Remote backend dispatch and colpali integration
+# ---------------------------------------------------------------------------
+
+class TestRemoteBackendDispatch:
+ """Verify RemoteVectorStore is selected and ColPaliModel integrates correctly."""
+
+ def _make_remote_model(self, url="http://localhost:18000"):
+ from unittest.mock import patch as _patch
+ with _patch("foretrieval.vector_db_server.client.VectorDBServerClient") as MockClient:
+ mock_client_inst = MagicMock()
+ mock_client_inst.open_collection.return_value = {
+ "opened": True, "backend": "qdrant", "created": True
+ }
+ MockClient.return_value = mock_client_inst
+ model = _make_model(
+ "remote",
+ storage_config={"url": url, "backend": "qdrant"},
+ )
+ return model, mock_client_inst
+
+ def test_remote_backend_selects_remote_store(self):
+ from foretrieval.vector_store.remote import RemoteVectorStore
+ model, _ = self._make_remote_model()
+ assert isinstance(model.vector_store, RemoteVectorStore)
+
+ def test_vector_store_is_open_remote(self):
+ """_vector_store_is_open returns True after open() is called."""
+ from foretrieval.colpali import ColPaliModel
+ model, mock_client = self._make_remote_model()
+ # Simulate open() having been called (set _opened flag)
+ model.vector_store._opened = True
+ assert model._vector_store_is_open() is True
+
+ def test_vector_store_is_not_open_before_open(self):
+ from foretrieval.colpali import ColPaliModel
+ model, _ = self._make_remote_model()
+ model.vector_store._opened = False
+ assert model._vector_store_is_open() is False
+
+ def test_set_processor_not_called_for_remote(self):
+ """set_processor is a local-only hook — must not be called for remote."""
+ from foretrieval.vector_store.remote import RemoteVectorStore
+ model, _ = self._make_remote_model()
+ assert isinstance(model.vector_store, RemoteVectorStore)
+ # RemoteVectorStore has no set_processor attribute
+ assert not hasattr(model.vector_store, "set_processor")
+
+ def test_export_index_writes_remote_backend(self, tmp_path):
+ """Remote mode stores bookkeeping on the server, not in a local dir."""
+ model, _ = self._make_remote_model()
+ model.index_root = str(tmp_path)
+ model.index_name = "test_export_remote"
+ model.doc_id_to_metadata = {1: {"title": "A"}}
+ model.doc_ids_to_file_names = {1: "a.pdf"}
+ model.embed_id_to_extra = {}
+ model.full_document_collection = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = 1
+ model.model_name = "vidore/colpali-v1.2"
+ model.storage_config = {"url": "http://localhost:18000", "backend": "qdrant"}
+
+ mock_store = MagicMock()
+ mock_store.supports_remote_bookkeeping.return_value = True
+ mock_store.load_bookkeeping.return_value = None
+ model.vector_store = mock_store
+
+ model._export_index()
+
+ # No local index directory should have been created.
+ assert not (tmp_path / "test_export_remote").exists()
+ # The bookkeeping blob must have been pushed to the server.
+ mock_store.export_bookkeeping.assert_called_once()
+ blob = mock_store.export_bookkeeping.call_args.args[0]
+ assert blob["index_config"]["storage_backend"] == "remote"
+ assert blob["doc_ids_to_file_names"] == {1: "a.pdf"}
+ # storage_config (with any credentials) must NOT be persisted in bookkeeping.
+ assert blob["index_config"].get("storage_config") is None
+
+ def test_export_index_strips_api_key(self, tmp_path):
+ """Connection config (incl. api_key) is never persisted in bookkeeping."""
+ model, _ = self._make_remote_model()
+ model.index_root = str(tmp_path)
+ model.index_name = "test_strip_key"
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.embed_id_to_extra = {}
+ model.full_document_collection = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.model_name = "vidore/colpali-v1.2"
+ model.storage_config = {
+ "url": "http://localhost:18000",
+ "backend": "qdrant",
+ "api_key": "supersecret",
+ }
+
+ mock_store = MagicMock()
+ mock_store.supports_remote_bookkeeping.return_value = True
+ mock_store.load_bookkeeping.return_value = None
+ model.vector_store = mock_store
+
+ model._export_index()
+
+ blob = mock_store.export_bookkeeping.call_args.args[0]
+ # No connection config / api_key anywhere in the persisted blob.
+ assert blob["index_config"].get("storage_config") is None
+ assert "supersecret" not in str(blob)
+
+ def test_indexed_embeddings_returns_empty_for_remote(self):
+ model, _ = self._make_remote_model()
+ assert model.indexed_embeddings == []
+
+ def test_embed_id_to_doc_id_returns_empty_for_remote(self):
+ model, _ = self._make_remote_model()
+ assert model.embed_id_to_doc_id == {}
diff --git a/tests/test_colpali_on_progress.py b/tests/test_colpali_on_progress.py
new file mode 100644
index 0000000..03ed644
--- /dev/null
+++ b/tests/test_colpali_on_progress.py
@@ -0,0 +1,159 @@
+"""Tests for the on_progress callback wiring in ColPaliModel.index().
+
+Avoids loading a real ColPali model — uses a minimal fake to verify that:
+ 1. The callback receives the expected stages in order
+ 2. Page events report sane page counts for a single-page PDF
+ 3. Exceptions raised in the callback are swallowed (best-effort delivery)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+
+# ---------------------------------------------------------------------------
+# Stub the ColPaliModel without instantiating the real one. We exercise the
+# index() function purely from the perspective of the on_progress callback.
+# This keeps the test fast (no model download / no CUDA needed).
+# ---------------------------------------------------------------------------
+
+class _FakeColPaliModel:
+ """A throwaway object exposing just enough of ColPaliModel.index's contract.
+
+ We hijack the bound method via __get__ to be able to call the real
+ ColPaliModel.index with `self` bound to this fake (so its branches that
+ only consult attributes work). The heavy lifting in add_to_index is
+ monkey-patched to a no-op that fires the standard page events.
+ """
+
+
+SAMPLE_PDF = Path(__file__).parent.parent / "sample_data" / "sample_doc.pdf"
+
+
+def test_on_progress_emits_expected_events(tmp_path, monkeypatch):
+ """index() must drive on_progress through start → file_start → file_done → all_done."""
+ if not SAMPLE_PDF.exists():
+ pytest.skip("sample PDF not available")
+
+ from foretrieval.colpali import ColPaliModel
+
+ events: list[dict] = []
+
+ def cb(evt: dict) -> None:
+ events.append(evt)
+
+ # Build a minimal fake "self" with the attributes index() reads
+ fake = _FakeColPaliModel()
+ fake.index_name = None
+ fake.index_root = str(tmp_path)
+ fake.storage_backend = "local"
+ fake.storage_config = {}
+ fake.full_document_collection = False
+ fake.vector_store = type("V", (), {
+ "open": lambda *a, **k: None,
+ "supports_remote_bookkeeping": lambda *a, **k: False,
+ })()
+ fake.processor = None
+ fake.doc_id_to_metadata = {}
+ fake.doc_ids = set()
+ fake.doc_ids_to_file_names = {}
+ fake.highest_doc_id = -1
+
+ def fake_add_to_index(self, item, store_collection_with_index, doc_id=None,
+ metadata=None, batch_size=1,
+ on_progress=None, _file_idx=0, _n_files=1):
+ # Simulate one page progress event so we can also verify pages bubble up
+ if on_progress is not None:
+ on_progress({
+ "stage": "page",
+ "file": Path(str(item)).name,
+ "file_idx": _file_idx,
+ "n_files": _n_files,
+ "page_idx": 0,
+ "n_pages": 1,
+ })
+ self.doc_ids_to_file_names[doc_id] = str(item)
+ self.highest_doc_id = max(self.highest_doc_id, doc_id)
+ self.doc_ids.add(doc_id)
+ return {}
+
+ fake.add_to_index = fake_add_to_index.__get__(fake, type(fake))
+ fake._export_index = lambda *a, **k: None
+ fake._vector_store_is_open = lambda: True
+
+ # Run the real index() bound to our fake
+ ColPaliModel.index(
+ fake,
+ input_path=SAMPLE_PDF.parent,
+ index_name="test_index_progress",
+ store_collection_with_index=False,
+ overwrite=True,
+ on_progress=cb,
+ )
+
+ stages = [e["stage"] for e in events]
+ assert "start" in stages
+ assert "file_start" in stages
+ assert "file_done" in stages
+ assert "all_done" in stages
+ # At least one page event must be present
+ assert any(s == "page" for s in stages)
+ # start must precede file_start which must precede file_done which must precede all_done
+ assert stages.index("start") < stages.index("file_start") < stages.index("file_done") < stages.index("all_done")
+
+
+def test_on_progress_exception_is_swallowed(tmp_path, monkeypatch):
+ """A callback that raises must not abort indexing."""
+ if not SAMPLE_PDF.exists():
+ pytest.skip("sample PDF not available")
+
+ from foretrieval.colpali import ColPaliModel
+
+ call_count = {"n": 0}
+
+ def bad_cb(evt):
+ call_count["n"] += 1
+ raise RuntimeError("intentional")
+
+ fake = _FakeColPaliModel()
+ fake.index_name = None
+ fake.index_root = str(tmp_path)
+ fake.storage_backend = "local"
+ fake.storage_config = {}
+ fake.full_document_collection = False
+ fake.vector_store = type("V", (), {
+ "open": lambda *a, **k: None,
+ "supports_remote_bookkeeping": lambda *a, **k: False,
+ })()
+ fake.processor = None
+ fake.doc_id_to_metadata = {}
+ fake.doc_ids = set()
+ fake.doc_ids_to_file_names = {}
+ fake.highest_doc_id = -1
+ fake._export_index = lambda *a, **k: None
+ fake._vector_store_is_open = lambda: True
+
+ def fake_add_to_index(self, item, store_collection_with_index, doc_id=None,
+ metadata=None, batch_size=1,
+ on_progress=None, _file_idx=0, _n_files=1):
+ self.doc_ids_to_file_names[doc_id] = str(item)
+ self.highest_doc_id = max(self.highest_doc_id, doc_id)
+ self.doc_ids.add(doc_id)
+ return {}
+
+ fake.add_to_index = fake_add_to_index.__get__(fake, type(fake))
+
+ # Should not raise even though every callback invocation throws
+ ColPaliModel.index(
+ fake,
+ input_path=SAMPLE_PDF.parent,
+ index_name="test_index_progress_bad_cb",
+ store_collection_with_index=False,
+ overwrite=True,
+ on_progress=bad_cb,
+ )
+
+ # And the callback was still invoked despite raising
+ assert call_count["n"] > 0
diff --git a/tests/test_colqwen.py b/tests/test_colqwen.py
deleted file mode 100644
index 6a68acc..0000000
--- a/tests/test_colqwen.py
+++ /dev/null
@@ -1,23 +0,0 @@
-from typing import Generator
-
-import pytest
-from colpali_engine.models import ColQwen2
-from colpali_engine.utils.torch_utils import get_torch_device, tear_down_torch
-
-from byaldi import RAGMultiModalModel
-from byaldi.colpali import ColPaliModel
-
-
-@pytest.fixture(scope="module")
-def colqwen_rag_model() -> Generator[RAGMultiModalModel, None, None]:
- device = get_torch_device("auto")
- print(f"Using device: {device}")
- yield RAGMultiModalModel.from_pretrained("vidore/colqwen2-v0.1", device=device)
- tear_down_torch()
-
-
-@pytest.mark.slow
-def test_load_colqwen_from_pretrained(colqwen_rag_model: RAGMultiModalModel):
- assert isinstance(colqwen_rag_model, RAGMultiModalModel)
- assert isinstance(colqwen_rag_model.model, ColPaliModel)
- assert isinstance(colqwen_rag_model.model.model, ColQwen2)
diff --git a/tests/test_e2e_rag.py b/tests/test_e2e_rag.py
index bba74dc..bdec46e 100644
--- a/tests/test_e2e_rag.py
+++ b/tests/test_e2e_rag.py
@@ -1,86 +1,77 @@
from pathlib import Path
from typing import Generator
-
+import shutil
import pytest
from colpali_engine.utils.torch_utils import get_torch_device, tear_down_torch
-from byaldi import RAGMultiModalModel
+from foretrieval import MultiModalRetrieverModel
-path_document_1 = Path("docs/attention.pdf")
-path_document_2 = Path("docs/attention_copy.pdf")
+path_document_1 = Path("sample_data/sample_pdf.pdf")
+path_document_2 = Path("sample_data/sample_multi_pdf.pdf")
+index_root = Path(".test_index")
@pytest.fixture(scope="function")
-def rag_model_from_pretrained() -> Generator[RAGMultiModalModel, None, None]:
+def rag_model_from_pretrained() -> Generator[MultiModalRetrieverModel, None, None]:
device = get_torch_device("auto")
print(f"Using device: {device}")
- yield RAGMultiModalModel.from_pretrained("vidore/colpali-v1.2", device=device)
+ yield MultiModalRetrieverModel.from_pretrained(
+ "vidore/colqwen2.5-v0.2", device=device
+ )
tear_down_torch()
@pytest.fixture(scope="function")
-def rag_model_from_index() -> Generator[RAGMultiModalModel, None, None]:
- yield RAGMultiModalModel.from_index("multi_doc_index")
+def rag_model_from_index() -> Generator[MultiModalRetrieverModel, None, None]:
+ if index_root.exists():
+ # delete the folder
+ shutil.rmtree(index_root)
+ yield MultiModalRetrieverModel.from_index(
+ "multi_doc_index", index_root=index_root.as_posix()
+ )
tear_down_torch()
@pytest.mark.slow
-def test_single_pdf(rag_model_from_pretrained: RAGMultiModalModel):
- if not Path("docs/attention.pdf").is_file():
- raise FileNotFoundError(
- f"Please download the PDF file from https://arxiv.org/pdf/1706.03762 and move it to {path_document_1}."
- )
-
+def test_single_pdf(rag_model_from_pretrained: MultiModalRetrieverModel):
# Index a single PDF
rag_model_from_pretrained.index(
- input_path="docs/attention.pdf",
- index_name="attention_index",
+ input_path=path_document_2,
+ index_name="sample_multipage_index",
store_collection_with_index=True,
overwrite=True,
)
# Test retrieval
queries = [
- "How does the positional encoding thing work?",
- "what's the BLEU score of this new strange method?",
+ "What is the answer we are looking for on the first page?",
+ "What is the answer we are looking for on the second page?",
+ "What is the answer we are looking for on the third page?",
]
+ expected_page = 1
for query in queries:
- results = rag_model_from_pretrained.search(query, k=3)
+ result = rag_model_from_pretrained.search(query, k=1)[0]
print(f"\nQuery: {query}")
- for result in results:
- print(
- f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
- )
+ print(
+ f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
+ )
- # Check if the expected page (6 for positional encoding) is in the top results
- if "positional encoding" in query.lower():
- assert any(
- r.page_num == 6 for r in results
- ), "Expected page 6 for positional encoding query"
+ # Check if the expected page is in the top results
- # Check if the expected pages (8 and 9 for BLEU score) are in the top results
- if "bleu score" in query.lower():
- assert any(
- r.page_num in [8, 9] for r in results
- ), "Expected pages 8 or 9 for BLEU score query"
+ assert result.page_num == expected_page, (
+ f"Expected page {expected_page} for this query. Got {result.page_num} instead"
+ )
+ expected_page += 1
-@pytest.mark.slow
-def test_multi_document(rag_model_from_pretrained: RAGMultiModalModel):
- if not Path("docs/attention.pdf").is_file():
- raise FileNotFoundError(
- f"Please download the PDF file from https://arxiv.org/pdf/1706.03762 and move it to {path_document_1}."
- )
- if not Path("docs/attention_copy.pdf").is_file():
- raise FileNotFoundError(
- f"Please download the PDF file from https://arxiv.org/pdf/1706.03762 and move it to {path_document_2}."
- )
+@pytest.mark.slow
+def test_multi_document(rag_model_from_pretrained: MultiModalRetrieverModel):
# Index a directory of documents
rag_model_from_pretrained.index(
- input_path="docs/",
+ input_path="sample_data/",
index_name="multi_doc_index",
store_collection_with_index=True,
overwrite=True,
@@ -88,66 +79,18 @@ def test_multi_document(rag_model_from_pretrained: RAGMultiModalModel):
# Test retrieval
queries = [
- "How does the positional encoding thing work?",
- "what's the BLEU score of this new strange method?",
+ "What is the numerical answer we are looking for in the sample xls file?",
+ "What is the numerical answer we are looking for on the third page of the sample multipage pdf file?",
]
for query in queries:
- results = rag_model_from_pretrained.search(query, k=5)
+ result = rag_model_from_pretrained.search(query, k=1)[0]
print(f"\nQuery: {query}")
- for result in results:
- print(
- f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
- )
-
- # Check if the expected page (6 for positional encoding) is in the top results
- if "positional encoding" in query.lower():
- assert any(
- r.page_num == 6 for r in results
- ), "Expected page 6 for positional encoding query"
-
- # Check if the expected pages (8 and 9 for BLEU score) are in the top results
- if "bleu score" in query.lower():
- assert any(
- r.page_num in [8, 9] for r in results
- ), "Expected pages 8 or 9 for BLEU score query"
-
-
-@pytest.mark.skip("This test should be made independent of the previous tests.")
-@pytest.mark.slow
-def test_add_to_index(rag_model_from_index: RAGMultiModalModel):
- # NOTE: This test should run after the test_multi_document test.
-
- # Add a new document to the index
- rag_model_from_index.add_to_index(
- input_item="docs/",
- store_collection_with_index=True,
- doc_id=[1002, 1003],
- metadata=[{"author": "John Doe", "year": 2023}] * 2,
- )
-
- # Test retrieval with the updated index
- queries = ["what's the BLEU score of this new strange method?"]
-
- for query in queries:
- results = rag_model_from_index.search(query, k=3)
+ print(
+ f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
+ )
- print(f"\nQuery: {query}")
- for result in results:
- print(
- f"Doc ID: {result.doc_id}, Page: {result.page_num}, Score: {result.score}"
- )
- print(f"Metadata: {result.metadata}")
-
- # Check if the expected page (6 for positional encoding) is in the top results
- if "positional encoding" in query.lower():
- assert any(
- r.page_num == 6 for r in results
- ), "Expected page 6 for positional encoding query"
-
- # Check if the expected pages (8 and 9 for BLEU score) are in the top results
- if "bleu score" in query.lower():
- assert any(
- r.page_num in [8, 9] for r in results
- ), "Expected pages 8 or 9 for BLEU score query"
+ # Check if the expected page (3) is in the top results
+ if "third" in query.lower():
+ assert result.page_num == 3, "Expected page 3 for multi-page pdf query"
diff --git a/tests/test_embedding_server.py b/tests/test_embedding_server.py
new file mode 100644
index 0000000..a4e0db2
--- /dev/null
+++ b/tests/test_embedding_server.py
@@ -0,0 +1,630 @@
+"""Unit tests for the embedding server package.
+
+All tests are offline — no real server, no GPU, no SSH required.
+HTTP and SSH interactions are fully mocked.
+"""
+
+import base64
+import io
+import json
+from unittest.mock import MagicMock, patch, call
+
+import httpx
+import pytest
+import torch
+from PIL import Image
+
+from foretrieval.embedding_server.config import EmbeddingServerConfig
+from foretrieval.embedding_server.client import (
+ EmbeddingServerClient,
+ ServerOOMError,
+ _pil_to_base64,
+)
+from foretrieval.embedding_server.manager import EmbeddingServerManager
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+_TEST_SERVER_URL = "http://localhost:8000"
+_TEST_MODEL = "athrael-soju/colqwen3.5-4.5B-v3"
+_TEST_SSH_HOST = "test-gpu-server"
+
+
+def _make_config(**kwargs) -> EmbeddingServerConfig:
+ defaults = dict(
+ url=_TEST_SERVER_URL,
+ model_name=_TEST_MODEL,
+ batch_size=4,
+ )
+ defaults.update(kwargs)
+ return EmbeddingServerConfig(**defaults)
+
+
+def _make_image(w=32, h=32) -> Image.Image:
+ return Image.new("RGB", (w, h), color=(128, 64, 32))
+
+
+def _fake_pooling_response(n_items: int, n_tokens: int = 10, dim: int = 128) -> dict:
+ """Build a fake vLLM /pooling response with n_items token-embed outputs."""
+ data = []
+ for i in range(n_items):
+ token_embeds = [[float(i) * 0.1] * dim for _ in range(n_tokens)]
+ data.append({"index": i, "object": "embedding", "data": token_embeds})
+ return {"object": "list", "data": data, "model": "test-model"}
+
+
+def _make_mock_response(status_code: int = 200, json_data: dict = None, text: str = "") -> MagicMock:
+ """Build a mock httpx.Response."""
+ resp = MagicMock(spec=httpx.Response)
+ resp.status_code = status_code
+ resp.json.return_value = json_data or {}
+ resp.text = text
+ return resp
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingServerConfig tests
+# ---------------------------------------------------------------------------
+
+class TestEmbeddingServerConfig:
+ def test_basic_construction(self):
+ cfg = _make_config()
+ assert cfg.url == _TEST_SERVER_URL
+ assert cfg.model_name == _TEST_MODEL
+ assert cfg.n_gpus == -1
+ assert cfg.batch_size == 4
+
+ def test_trailing_slash_stripped(self):
+ cfg = _make_config(url="http://localhost:8000/")
+ assert cfg.url == "http://localhost:8000"
+
+ def test_n_gpus_zero_invalid(self):
+ with pytest.raises(Exception):
+ _make_config(n_gpus=0)
+
+ def test_n_gpus_negative_minus_one_valid(self):
+ cfg = _make_config(n_gpus=-1)
+ assert cfg.n_gpus == -1
+
+ def test_n_gpus_positive_valid(self):
+ cfg = _make_config(n_gpus=2)
+ assert cfg.n_gpus == 2
+
+ def test_auto_deploy_without_ssh_host_raises(self):
+ with pytest.raises(Exception, match="ssh_host"):
+ _make_config(auto_deploy=True, ssh_host=None)
+
+ def test_auto_deploy_with_ssh_host_ok(self):
+ cfg = _make_config(auto_deploy=True, ssh_host=_TEST_SSH_HOST)
+ assert cfg.auto_deploy is True
+
+ def test_from_dict(self):
+ d = {"url": "http://host:9000", "model_name": "athrael-soju/colqwen3.5-4.5B-v3"}
+ cfg = EmbeddingServerConfig.from_dict(d)
+ assert cfg.url == "http://host:9000"
+ assert cfg.port == 8000 # default
+
+ # --- model_name validator ---
+
+ def test_incompatible_model_raises(self):
+ with pytest.raises(ValueError, match="colqwen3"):
+ _make_config(model_name="vidore/colpali-v1.2")
+
+ def test_colqwen2_model_raises(self):
+ with pytest.raises(ValueError, match="colqwen3"):
+ _make_config(model_name="vidore/colqwen2-v1.0")
+
+ def test_colqwen3_model_accepted(self):
+ cfg = _make_config(model_name="athrael-soju/colqwen3.5-4.5B-v3")
+ assert "colqwen3" in cfg.model_name.lower()
+
+ def test_colqwen3_variant_accepted(self):
+ cfg = _make_config(model_name="vidore/colqwen3-v1.0")
+ assert cfg.model_name == "vidore/colqwen3-v1.0"
+
+ # --- auth / SSL fields ---
+
+ def test_api_key_default_none(self):
+ cfg = _make_config()
+ assert cfg.api_key is None
+
+ def test_api_key_set(self):
+ cfg = _make_config(api_key="secret-token")
+ assert cfg.api_key == "secret-token"
+
+ def test_verify_ssl_default_true(self):
+ cfg = _make_config()
+ assert cfg.verify_ssl is True
+
+ def test_verify_ssl_false(self):
+ cfg = _make_config(verify_ssl=False)
+ assert cfg.verify_ssl is False
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingServerClient — request format
+# ---------------------------------------------------------------------------
+
+class TestEmbeddingServerClientRequestFormat:
+ def _make_client(self, **kwargs) -> EmbeddingServerClient:
+ return EmbeddingServerClient(_make_config(**kwargs))
+
+ def test_embed_images_sends_correct_payload(self):
+ client = self._make_client()
+ images = [_make_image()]
+ fake_resp = _fake_pooling_response(1)
+
+ mock_response = _make_mock_response(200, fake_resp)
+
+ with patch.object(client._client, "post", return_value=mock_response) as mock_post:
+ client.embed_images(images)
+
+ mock_post.assert_called_once()
+ _, kwargs = mock_post.call_args
+ payload = kwargs["json"]
+
+ assert payload["model"] == "athrael-soju/colqwen3.5-4.5B-v3"
+ # vLLM >=0.19.0: images sent via messages array (PoolingChatRequest)
+ assert "messages" in payload
+ assert len(payload["messages"]) == 1
+ content = payload["messages"][0]["content"]
+ assert len(content) == 1
+ assert content[0]["type"] == "image_url"
+ assert content[0]["image_url"]["url"].startswith("data:image/png;base64,")
+
+ def test_embed_query_sends_text_payload(self):
+ client = self._make_client()
+ fake_resp = _fake_pooling_response(1, n_tokens=5)
+
+ mock_response = _make_mock_response(200, fake_resp)
+
+ with patch.object(client._client, "post", return_value=mock_response) as mock_post:
+ client.embed_query("What is the speed?")
+
+ _, kwargs = mock_post.call_args
+ payload = kwargs["json"]
+ assert payload["input"] == "What is the speed?"
+ assert "task" not in payload
+
+ def test_embed_images_uses_correct_endpoint(self):
+ client = self._make_client()
+ mock_response = _make_mock_response(200, _fake_pooling_response(1))
+
+ with patch.object(client._client, "post", return_value=mock_response) as mock_post:
+ client.embed_images([_make_image()])
+
+ url_called = mock_post.call_args[0][0]
+ assert url_called == f"{_TEST_SERVER_URL}/pooling"
+
+ def test_auth_header_sent_when_api_key_set(self):
+ """api_key set → Authorization: Bearer header in httpx.Client headers."""
+ cfg = _make_config(api_key="my-secret")
+ client = EmbeddingServerClient(cfg)
+ assert client._client.headers.get("authorization") == "Bearer my-secret"
+
+ def test_no_auth_header_when_api_key_none(self):
+ """api_key=None → no Authorization header."""
+ cfg = _make_config()
+ client = EmbeddingServerClient(cfg)
+ assert "authorization" not in client._client.headers
+
+ def test_ssl_verify_true_by_default(self):
+ """verify_ssl=True → httpx.Client verifies SSL."""
+ cfg = _make_config()
+ client = EmbeddingServerClient(cfg)
+ # httpx.Client stores ssl_context; verify=True means it's not disabled
+ # We check the config was respected via the verify attr on the transport
+ assert client.config.verify_ssl is True
+
+ def test_ssl_verify_false_passed_to_client(self):
+ """verify_ssl=False → httpx.Client created with verify=False."""
+ cfg = _make_config(verify_ssl=False)
+ # Should not raise; httpx accepts verify=False
+ client = EmbeddingServerClient(cfg)
+ assert client.config.verify_ssl is False
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingServerClient — response parsing
+# ---------------------------------------------------------------------------
+
+class TestEmbeddingServerClientResponseParsing:
+ def _client_with_response(self, response_dict: dict) -> EmbeddingServerClient:
+ client = EmbeddingServerClient(_make_config())
+ mock_response = _make_mock_response(200, response_dict)
+ client._client = MagicMock()
+ client._client.post.return_value = mock_response
+ return client
+
+ def test_embed_images_returns_correct_number_of_tensors(self):
+ n = 3
+ client = self._client_with_response(_fake_pooling_response(n))
+ result = client.embed_images([_make_image()] * n)
+ assert len(result) == n
+
+ def test_embed_images_tensor_shape(self):
+ n_tokens, dim = 12, 128
+ client = self._client_with_response(_fake_pooling_response(1, n_tokens, dim))
+ result = client.embed_images([_make_image()])
+ assert result[0].shape == (n_tokens, dim)
+
+ def test_embed_images_tensor_dtype_float32(self):
+ client = self._client_with_response(_fake_pooling_response(1))
+ result = client.embed_images([_make_image()])
+ assert result[0].dtype == torch.float32
+
+ def test_embed_query_returns_single_tensor(self):
+ client = self._client_with_response(_fake_pooling_response(1, n_tokens=7, dim=128))
+ result = client.embed_query("hello")
+ assert len(result) == 1
+ assert result[0].shape == (7, 128)
+
+ def test_embed_empty_images_returns_empty(self):
+ client = EmbeddingServerClient(_make_config())
+ result = client.embed_images([])
+ assert result == []
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingServerClient — OOM retry
+# ---------------------------------------------------------------------------
+
+class TestEmbeddingServerClientOOMRetry:
+ def test_oom_halves_batch_size_and_retries(self):
+ """Server OOM on batch_size=4 → retries with 2 → succeeds."""
+ client = EmbeddingServerClient(_make_config(batch_size=4))
+
+ call_count = {"n": 0}
+ first_call = {"done": False}
+
+ def fake_post(url, json=None, **kwargs):
+ call_count["n"] += 1
+ if not first_call["done"]:
+ first_call["done"] = True
+ return _make_mock_response(500, text="CUDA out of memory trying to allocate tensor")
+ return _make_mock_response(200, _fake_pooling_response(1))
+
+ client._client = MagicMock()
+ client._client.post.side_effect = fake_post
+
+ images = [_make_image()] * 4
+ result = client.embed_images(images)
+ assert len(result) == 4
+ assert call_count["n"] > 1
+
+ def test_oom_at_batch_size_1_raises(self):
+ """OOM even at batch_size=1 → raises ServerOOMError."""
+ client = EmbeddingServerClient(_make_config(batch_size=1))
+
+ def fake_post(url, json=None, **kwargs):
+ return _make_mock_response(500, text="CUDA out of memory")
+
+ client._client = MagicMock()
+ client._client.post.side_effect = fake_post
+
+ with pytest.raises(ServerOOMError):
+ client.embed_images([_make_image()])
+
+ def test_non_oom_500_raises_runtime_error(self):
+ client = EmbeddingServerClient(_make_config())
+
+ def fake_post(url, json=None, **kwargs):
+ return _make_mock_response(500, text="Internal server error: model not loaded")
+
+ client._client = MagicMock()
+ client._client.post.side_effect = fake_post
+
+ with pytest.raises(RuntimeError, match="HTTP 500"):
+ client.embed_images([_make_image()])
+
+ def test_connection_error_raises(self):
+ client = EmbeddingServerClient(_make_config())
+ client._client = MagicMock()
+ client._client.post.side_effect = httpx.ConnectError("refused")
+
+ with pytest.raises(ConnectionError):
+ client.embed_images([_make_image()])
+
+ def test_timeout_raises(self):
+ client = EmbeddingServerClient(_make_config())
+ client._client = MagicMock()
+ client._client.post.side_effect = httpx.TimeoutException("timed out")
+
+ with pytest.raises(TimeoutError):
+ client.embed_images([_make_image()])
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingServerClient — health check
+# ---------------------------------------------------------------------------
+
+class TestHealthCheck:
+ def test_health_check_true_on_200(self):
+ client = EmbeddingServerClient(_make_config())
+ mock_resp = _make_mock_response(200)
+ client._client = MagicMock()
+ client._client.get.return_value = mock_resp
+ assert client.health_check() is True
+
+ def test_health_check_false_on_500(self):
+ client = EmbeddingServerClient(_make_config())
+ mock_resp = _make_mock_response(500)
+ client._client = MagicMock()
+ client._client.get.return_value = mock_resp
+ assert client.health_check() is False
+
+ def test_health_check_false_on_connection_error(self):
+ client = EmbeddingServerClient(_make_config())
+ client._client = MagicMock()
+ client._client.get.side_effect = httpx.ConnectError("refused")
+ assert client.health_check() is False
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingServerManager — deploy logic (SSH mocked)
+# ---------------------------------------------------------------------------
+
+class TestEmbeddingServerManager:
+ def _make_manager(self, **kwargs) -> EmbeddingServerManager:
+ defaults = dict(auto_deploy=True, ssh_host=_TEST_SSH_HOST, n_gpus=-1)
+ defaults.update(kwargs)
+ cfg = _make_config(**defaults)
+ mgr = EmbeddingServerManager(cfg)
+ # Stub the remote-home resolution so tests that mock _run_remote
+ # don't accidentally trigger a real SSH connection through
+ # _remote_home() -> sftp.normalize('.').
+ mgr._cached_home = "/home/testuser"
+ return mgr
+
+ def _mock_ssh(self, manager: EmbeddingServerManager, remote_outputs: dict):
+ def fake_run(cmd, on_line=None):
+ for key, val in remote_outputs.items():
+ if key in cmd:
+ return val
+ return ("", "")
+ manager._run_remote = MagicMock(side_effect=fake_run)
+ return manager
+
+ def test_deploy_when_no_metadata(self):
+ """No metadata file → deploy from scratch."""
+ mgr = self._make_manager()
+ self._mock_ssh(mgr, {
+ "deployment.json": ("__MISSING__", ""),
+ "nvidia-smi": ("2\n", ""),
+ "docker pull": ("", ""),
+ "docker run": ("abc123\n", ""),
+ })
+
+ mgr.ensure_deployed()
+
+ calls = [str(c) for c in mgr._run_remote.call_args_list]
+ assert any("docker pull" in c for c in calls)
+ assert any("docker run" in c for c in calls)
+
+ def test_deploy_skipped_when_container_running(self):
+ """Metadata present + container running → no redeploy."""
+ mgr = self._make_manager()
+ metadata = json.dumps({
+ "model_name": "athrael-soju/colqwen3.5-4.5B-v3",
+ "container_name": "foretrieval_embedding_server",
+ "port": 8000,
+ })
+ self._mock_ssh(mgr, {
+ "deployment.json": (metadata, ""),
+ "docker inspect": ("true\n", ""),
+ })
+
+ mgr.ensure_deployed()
+
+ calls = [str(c) for c in mgr._run_remote.call_args_list]
+ assert not any("docker run" in c for c in calls)
+
+ def test_redeploy_when_container_not_running(self):
+ """Metadata present but container stopped → redeploy."""
+ mgr = self._make_manager()
+ metadata = json.dumps({
+ "model_name": "athrael-soju/colqwen3.5-4.5B-v3",
+ "container_name": "foretrieval_embedding_server",
+ "port": 8000,
+ })
+ self._mock_ssh(mgr, {
+ "deployment.json": (metadata, ""),
+ "docker inspect": ("false\n", ""),
+ "nvidia-smi": ("2\n", ""),
+ "docker pull": ("", ""),
+ "docker run": ("abc123\n", ""),
+ })
+
+ mgr.ensure_deployed()
+
+ calls = [str(c) for c in mgr._run_remote.call_args_list]
+ assert any("docker run" in c for c in calls)
+
+ def test_n_gpus_minus1_uses_all_detected(self):
+ """n_gpus=-1 → detect GPU count via nvidia-smi."""
+ mgr = self._make_manager(n_gpus=-1)
+ self._mock_ssh(mgr, {
+ "deployment.json": ("__MISSING__", ""),
+ "nvidia-smi": ("2\n", ""),
+ "docker pull": ("", ""),
+ "docker run": ("abc123\n", ""),
+ })
+
+ mgr.ensure_deployed()
+
+ docker_run_calls = [
+ str(c) for c in mgr._run_remote.call_args_list
+ if "docker run" in str(c)
+ ]
+ assert docker_run_calls
+ assert "tensor-parallel-size 2" in docker_run_calls[0]
+
+ def test_n_gpus_explicit_used_directly(self):
+ """n_gpus=1 → skip nvidia-smi, use 1 directly."""
+ mgr = self._make_manager(n_gpus=1)
+ self._mock_ssh(mgr, {
+ "deployment.json": ("__MISSING__", ""),
+ "docker pull": ("", ""),
+ "docker run": ("abc123\n", ""),
+ })
+
+ mgr.ensure_deployed()
+
+ docker_run_calls = [
+ str(c) for c in mgr._run_remote.call_args_list
+ if "docker run" in str(c)
+ ]
+ assert "tensor-parallel-size 1" in docker_run_calls[0]
+ nvidia_calls = [
+ c for c in mgr._run_remote.call_args_list
+ if "nvidia-smi" in str(c)
+ ]
+ assert not nvidia_calls
+
+ def test_stop_removes_container_and_metadata(self):
+ mgr = self._make_manager()
+ mgr._run_remote = MagicMock(return_value=("", ""))
+ mgr.stop()
+ calls = [str(c) for c in mgr._run_remote.call_args_list]
+ assert any("docker stop" in c for c in calls)
+ assert any("docker rm" in c for c in calls)
+ assert any("deployment.json" in c for c in calls)
+
+ def test_missing_paramiko_raises_import_error(self):
+ mgr = self._make_manager()
+ with patch.dict("sys.modules", {"paramiko": None}):
+ with pytest.raises(ImportError, match="paramiko"):
+ mgr.ensure_deployed()
+
+ def test_redeploy_calls_deploy(self):
+ mgr = self._make_manager()
+ mgr._deploy = MagicMock()
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.redeploy()
+ mgr._deploy.assert_called_once()
+
+ def test_redeploy_forwards_on_line_callback(self):
+ mgr = self._make_manager()
+ mgr._deploy = MagicMock()
+ cb = MagicMock()
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.redeploy(on_line=cb)
+ assert mgr._deploy.call_args.kwargs.get("on_line") is cb
+
+ def test_run_remote_streams_lines_when_callback_given(self):
+ """When on_line is provided, stdout is read line-by-line."""
+ mgr = self._make_manager()
+
+ class _FakeChannel:
+ def recv_exit_status(self):
+ return 0
+
+ class _FakeStdout:
+ def __init__(self, lines):
+ self._lines = list(lines)
+ self.channel = _FakeChannel()
+
+ def readline(self):
+ if self._lines:
+ return self._lines.pop(0)
+ return ""
+
+ def read(self):
+ return b""
+
+ fake_stdout = _FakeStdout(["pulling…\n", "complete\n"])
+ fake_stderr = MagicMock()
+ fake_stderr.read.return_value = b""
+ fake_ssh = MagicMock()
+ fake_ssh.exec_command.return_value = (None, fake_stdout, fake_stderr)
+ mgr._get_ssh = MagicMock(return_value=fake_ssh)
+
+ captured: list[str] = []
+ mgr._run_remote("docker pull vllm/vllm-openai", on_line=captured.append)
+ assert captured == ["pulling…", "complete"]
+
+ def test_get_ssh_delegates_to_open_ssh_client(self):
+ """_get_ssh must call foretrieval.ssh_utils.open_ssh_client.
+
+ Regression test for the bug where paramiko was given a raw alias
+ and DNS-failed for entries that only existed in ~/.ssh/config.
+ """
+ mgr = self._make_manager(ssh_user="alice", ssh_key_path="/k")
+ fake_client = MagicMock()
+ with patch("foretrieval.ssh_utils.open_ssh_client",
+ return_value=fake_client) as mock_open:
+ result = mgr._get_ssh()
+ mock_open.assert_called_once_with(
+ ssh_host=_TEST_SSH_HOST,
+ ssh_user="alice",
+ ssh_key_path="/k",
+ )
+ assert result is fake_client
+
+ def test_remote_home_uses_sftp_normalize(self):
+ mgr = self._make_manager()
+ mgr._cached_home = None # reset the stub from _make_manager
+
+ fake_sftp = MagicMock()
+ fake_sftp.normalize.return_value = "/home/alice"
+ fake_ssh = MagicMock()
+ fake_ssh.open_sftp.return_value = fake_sftp
+ mgr._get_ssh = MagicMock(return_value=fake_ssh)
+
+ assert mgr._remote_home() == "/home/alice"
+ fake_sftp.normalize.assert_called_once_with(".")
+
+ def test_metadata_path_is_absolute(self):
+ mgr = self._make_manager() # _cached_home = "/home/testuser"
+ path = mgr._metadata_path()
+ assert path.startswith("/home/testuser/")
+ assert "~" not in path
+ assert path.endswith("deployment.json")
+
+
+# ---------------------------------------------------------------------------
+# Quantization config (local mode, no GPU needed — just validates ctor)
+# ---------------------------------------------------------------------------
+
+class TestQuantizationConfig:
+ def test_colpali_model_accepts_quantization_params(self):
+ """ColPaliModel.__init__ signature accepts load_in_4bit/8bit — smoke test."""
+ import inspect
+ from foretrieval.colpali import ColPaliModel
+ sig = inspect.signature(ColPaliModel.__init__)
+ assert "load_in_4bit" in sig.parameters
+ assert "load_in_8bit" in sig.parameters
+ assert "bnb_4bit_quant_type" in sig.parameters
+ assert "bnb_4bit_compute_dtype" in sig.parameters
+
+ def test_retriever_model_accepts_quantization_params(self):
+ """MultiModalRetrieverModel.from_pretrained accepts quantization params."""
+ import inspect
+ from foretrieval.retriever import MultiModalRetrieverModel
+ sig = inspect.signature(MultiModalRetrieverModel.from_pretrained)
+ assert "load_in_4bit" in sig.parameters
+ assert "load_in_8bit" in sig.parameters
+
+ def test_quantization_without_bitsandbytes_raises(self):
+ """load_in_4bit=True with missing bitsandbytes → ImportError at load time."""
+ from foretrieval.colpali import ColPaliModel
+ with patch.dict("sys.modules", {"transformers.utils.bitsandbytes": None}):
+ with patch("builtins.__import__", side_effect=lambda name, *a, **kw: (
+ (_ for _ in ()).throw(ImportError("No module named 'bitsandbytes'"))
+ if name == "bitsandbytes" else __import__(name, *a, **kw)
+ )):
+ pass # import-level mock; actual test is integration-only without GPU
+
+
+# ---------------------------------------------------------------------------
+# Utility
+# ---------------------------------------------------------------------------
+
+class TestPilToBase64:
+ def test_returns_valid_base64_png(self):
+ img = _make_image()
+ b64 = _pil_to_base64(img)
+ raw = base64.b64decode(b64)
+ assert raw[:4] == b"\x89PNG"
diff --git a/tests/test_embedding_server_integration.py b/tests/test_embedding_server_integration.py
new file mode 100644
index 0000000..72769d4
--- /dev/null
+++ b/tests/test_embedding_server_integration.py
@@ -0,0 +1,279 @@
+"""Integration tests for the remote embedding server.
+
+These tests require a live vLLM server and are skipped by default.
+
+Environment variables
+---------------------
+FORETRIEVAL_TEST_SERVER : str
+ Base URL of the vLLM server to test against, e.g. "http://gpu-server:8000".
+ Tests are skipped when this variable is not set.
+FORETRIEVAL_TEST_SSH_HOST : str
+ SSH hostname for deployment tests (auto_deploy path), e.g. "gpu-server".
+ Only needed for the deploy/stop tests.
+FORETRIEVAL_TEST_MODEL : str
+ HuggingFace model ID served by the server.
+ Defaults to "athrael-soju/colqwen3.5-4.5B-v3".
+
+Usage
+-----
+ # Run all integration tests against a remote GPU server:
+ FORETRIEVAL_TEST_SERVER=http://:8000 \
+ FORETRIEVAL_TEST_SSH_HOST= \
+ uv run pytest tests/test_embedding_server_integration.py -m "integration" -v
+"""
+
+import os
+
+import pytest
+import torch
+from PIL import Image
+
+from foretrieval.embedding_server.client import EmbeddingServerClient
+from foretrieval.embedding_server.config import EmbeddingServerConfig
+from foretrieval.embedding_server.manager import EmbeddingServerManager
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+_SERVER_URL = os.environ.get("FORETRIEVAL_TEST_SERVER", "")
+_SSH_HOST = os.environ.get("FORETRIEVAL_TEST_SSH_HOST", "")
+_MODEL = os.environ.get(
+ "FORETRIEVAL_TEST_MODEL", "athrael-soju/colqwen3.5-4.5B-v3"
+)
+
+requires_server = pytest.mark.skipif(
+ not _SERVER_URL,
+ reason="Set FORETRIEVAL_TEST_SERVER=http://: to run integration tests",
+)
+requires_ssh = pytest.mark.skipif(
+ not _SSH_HOST,
+ reason="Set FORETRIEVAL_TEST_SSH_HOST= to run deploy integration tests",
+)
+
+
+@pytest.fixture(scope="module")
+def server_config() -> EmbeddingServerConfig:
+ return EmbeddingServerConfig(
+ url=_SERVER_URL,
+ model_name=_MODEL,
+ batch_size=2,
+ request_timeout=180,
+ )
+
+
+@pytest.fixture(scope="module")
+def client(server_config) -> EmbeddingServerClient:
+ return EmbeddingServerClient(server_config)
+
+
+def _make_image(w: int = 64, h: int = 64) -> Image.Image:
+ return Image.new("RGB", (w, h), color=(100, 150, 200))
+
+
+# ---------------------------------------------------------------------------
+# Health check
+# ---------------------------------------------------------------------------
+
+@requires_server
+@pytest.mark.integration
+def test_server_health(client):
+ assert client.health_check(), (
+ f"Server at {_SERVER_URL} is not healthy — is vLLM running?"
+ )
+
+
+# ---------------------------------------------------------------------------
+# Image embedding
+# ---------------------------------------------------------------------------
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_single_image_returns_tensor(client):
+ imgs = [_make_image()]
+ result = client.embed_images(imgs)
+ assert len(result) == 1
+ assert isinstance(result[0], torch.Tensor)
+ assert result[0].ndim == 2 # [n_tokens, dim]
+ assert result[0].shape[1] == 320 # ColQwen3.5 embed dim
+
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_multiple_images_returns_correct_count(client):
+ n = 3
+ imgs = [_make_image()] * n
+ result = client.embed_images(imgs)
+ assert len(result) == n
+
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_images_tensors_are_on_cpu(client):
+ result = client.embed_images([_make_image()])
+ assert result[0].device.type == "cpu"
+
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_images_tensors_differ_across_images(client):
+ """Different images should yield different embeddings."""
+ img_a = Image.new("RGB", (64, 64), color=(255, 0, 0))
+ img_b = Image.new("RGB", (64, 64), color=(0, 0, 255))
+ emb_a, emb_b = client.embed_images([img_a, img_b])
+ assert not torch.allclose(emb_a, emb_b)
+
+
+# ---------------------------------------------------------------------------
+# Query embedding
+# ---------------------------------------------------------------------------
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_query_returns_single_tensor(client):
+ result = client.embed_query("What is the maximum speed?")
+ assert len(result) == 1
+ assert isinstance(result[0], torch.Tensor)
+ assert result[0].ndim == 2
+ assert result[0].shape[1] == 320
+
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_query_differs_from_image_embedding(client):
+ """Query and image embeddings should be in the same space but differ."""
+ q_emb = client.embed_query("What is the airspeed?")[0]
+ img_emb = client.embed_images([_make_image()])[0]
+ # Both are [n_tokens, 320] — shapes may differ but dtypes should match
+ assert q_emb.dtype == img_emb.dtype
+
+
+# ---------------------------------------------------------------------------
+# Batch size / OOM retry (live server, small batches)
+# ---------------------------------------------------------------------------
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_images_batch_larger_than_one(client):
+ """Server should handle multiple images per request without OOM on RTX4090."""
+ imgs = [_make_image(128, 128)] * 4
+ result = client.embed_images(imgs)
+ assert len(result) == 4
+
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_embed_empty_list_returns_empty(client):
+ result = client.embed_images([])
+ assert result == []
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: ColPaliModel with remote client
+# ---------------------------------------------------------------------------
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_colpali_model_remote_embed_images(server_config, tmp_path):
+ """ColPaliModel loads processor only and uses remote client for embeddings."""
+ from foretrieval.colpali import ColPaliModel
+
+ model = ColPaliModel.from_pretrained(
+ pretrained_model_name_or_path=_MODEL,
+ index_root=str(tmp_path),
+ embedding_server=server_config,
+ storage_qdrant=False, # avoid qdrant dep in integration test env
+ )
+ assert model._remote_client is not None
+ assert model.model is None # no local weights loaded
+
+ # Index a small synthetic image
+ img = _make_image(64, 64)
+ model.index_name = "test_remote"
+ model.highest_doc_id = -1
+ model.full_document_collection = False
+ model.resize_stored_images = False
+ model.max_image_width = None
+ model.max_image_height = None
+
+ model._add_to_index(
+ images=[img],
+ store_collection_with_index=False,
+ doc_id=0,
+ page_ids=[1],
+ )
+
+ # Should have one embedding stored
+ if model.storage_qdrant:
+ # qdrant path — count points
+ pass # point insertion verified by no exception
+ else:
+ assert len(model.indexed_embeddings) == 1
+ assert model.indexed_embeddings[0].shape[1] == 320 # ColQwen3.5 dim
+
+
+@requires_server
+@pytest.mark.integration
+@pytest.mark.slow
+def test_colpali_model_remote_encode_query(server_config, tmp_path):
+ """_encode_search_query uses remote client in remote mode."""
+ from foretrieval.colpali import ColPaliModel
+
+ model = ColPaliModel.from_pretrained(
+ pretrained_model_name_or_path=_MODEL,
+ index_root=str(tmp_path),
+ embedding_server=server_config,
+ )
+ qs = model._encode_search_query("What is the fuel capacity?")
+ assert len(qs) == 1
+ assert isinstance(qs[0], torch.Tensor)
+ assert qs[0].ndim == 2
+
+
+# ---------------------------------------------------------------------------
+# Deployment manager (requires SSH access)
+# ---------------------------------------------------------------------------
+
+@requires_server
+@requires_ssh
+@pytest.mark.integration
+@pytest.mark.slow
+def test_manager_health_check_against_running_server():
+ """If server is already running, health check returns True."""
+ cfg = EmbeddingServerConfig(
+ url=_SERVER_URL,
+ model_name=_MODEL,
+ ssh_host=_SSH_HOST,
+ )
+ client = EmbeddingServerClient(cfg)
+ assert client.health_check()
+
+
+@requires_ssh
+@pytest.mark.integration
+@pytest.mark.slow
+def test_manager_ensure_deployed_idempotent():
+ """ensure_deployed called twice should not start a second container."""
+ cfg = EmbeddingServerConfig(
+ url=_SERVER_URL or f"http://{_SSH_HOST}:8000",
+ model_name=_MODEL,
+ auto_deploy=True,
+ ssh_host=_SSH_HOST,
+ n_gpus=-1,
+ )
+ mgr = EmbeddingServerManager(cfg)
+ # First call: may deploy or detect already running
+ mgr.ensure_deployed()
+ # Second call: container already running → no-op
+ mgr.ensure_deployed()
+ # Verify container is running
+ assert mgr._is_container_running()
diff --git a/tests/test_init_lazy_import.py b/tests/test_init_lazy_import.py
new file mode 100644
index 0000000..56e998b
--- /dev/null
+++ b/tests/test_init_lazy_import.py
@@ -0,0 +1,105 @@
+"""Tests for foretrieval.__init__ — specifically the lazy-import guarantee.
+
+The vector-DB server is started by Python importing
+``foretrieval.vector_db_server.server_main`` which triggers the
+top-level ``foretrieval/__init__.py``. If that __init__ eagerly
+imports ``colpali.py``, it pulls in ``colpali_engine``, ``transformers``
+and ``torch._dynamo``, which crashes on the CPU-only Docker image with:
+
+ AssertionError: Artifact of type=precompile already registered
+ in mega-cache artifact factory
+
+The fix is a lazy ``__getattr__``-based __init__ so the heavy ML stack
+is only loaded when the caller explicitly requests it.
+"""
+
+from __future__ import annotations
+
+import importlib
+import sys
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _restore_sys_modules():
+ """Snapshot sys.modules before each test and restore it afterward.
+
+ Each test in this file deliberately purges all foretrieval.* entries from
+ sys.modules to verify lazy-import semantics. Without this fixture the purge
+ is permanent for the rest of the pytest session, causing later tests to
+ receive split module identities (OLD class object vs NEW class object) that
+ break isinstance() checks, patch() targets, and module-global reads.
+ """
+ snapshot = {k: v for k, v in sys.modules.items()}
+ yield
+ # Remove any modules that were added during the test (not in snapshot)
+ for k in [k for k in sys.modules if k not in snapshot]:
+ del sys.modules[k]
+ # Restore all snapshot entries, overwriting any re-imported variants
+ sys.modules.update(snapshot)
+
+
+def test_foretrieval_init_does_not_import_colpali():
+ """Importing ``foretrieval`` must NOT pull in colpali as a side-effect.
+
+ This is the regression guard for the vector-DB server startup crash:
+ the server only needs ``foretrieval.vector_db_server.*`` and must be
+ able to start without torch / colpali_engine being importable.
+ """
+ # Remove any already-cached foretrieval / colpali modules so the
+ # test is independent of import order within the test session.
+ to_remove = [k for k in sys.modules if k.startswith("foretrieval") or "colpali" in k]
+ for k in to_remove:
+ sys.modules.pop(k, None)
+
+ colpali_before = {k for k in sys.modules if "colpali" in k}
+
+ import foretrieval # noqa: F401
+
+ colpali_after = {k for k in sys.modules if "colpali" in k}
+
+ new_colpali = colpali_after - colpali_before
+ assert not new_colpali, (
+ f"Importing 'foretrieval' pulled in colpali modules: {new_colpali}. "
+ "The __init__.py must use lazy imports so the vector-DB server "
+ "can start without torch/colpali_engine being available."
+ )
+
+
+def test_foretrieval_multimodal_retriever_model_accessible():
+ """MultiModalRetrieverModel must still be reachable via the package."""
+ # Re-import fresh
+ to_remove = [k for k in sys.modules if k.startswith("foretrieval")]
+ for k in to_remove:
+ sys.modules.pop(k, None)
+
+ import foretrieval
+ cls = foretrieval.MultiModalRetrieverModel
+ assert cls is not None
+ # Confirm it is the real class (not a stub)
+ from foretrieval.retriever import MultiModalRetrieverModel as Direct
+ assert cls is Direct
+
+
+def test_foretrieval_ai_metadata_provider_accessible():
+ """ai_metadata_provider_factory must still be reachable via the package."""
+ to_remove = [k for k in sys.modules if k.startswith("foretrieval")]
+ for k in to_remove:
+ sys.modules.pop(k, None)
+
+ import foretrieval
+ fn = foretrieval.ai_metadata_provider_factory
+ assert callable(fn)
+
+
+def test_foretrieval_unknown_attr_raises():
+ """Accessing a non-existent attribute must raise AttributeError, not hang."""
+ to_remove = [k for k in sys.modules if k.startswith("foretrieval")]
+ for k in to_remove:
+ sys.modules.pop(k, None)
+
+ import foretrieval
+ import pytest
+ with pytest.raises(AttributeError, match="foretrieval"):
+ _ = foretrieval.this_does_not_exist
diff --git a/tests/test_metadata_ai.py b/tests/test_metadata_ai.py
new file mode 100644
index 0000000..b496049
--- /dev/null
+++ b/tests/test_metadata_ai.py
@@ -0,0 +1,273 @@
+"""
+Tests for FORetrieval metadata generation WITH an AI provider.
+
+These tests exercise the ``_provider`` callable returned by
+``ai_metadata_provider_factory(ai_cfg)`` when a real LLM is configured.
+They require at least one of the following environment variables:
+
+ OPENROUTER_API_KEY (preferred — fast, high rate limits)
+ Model: mistralai/mistral-small-3.2-24b-instruct
+ OPENAI_API_KEY Model: gpt-4o-mini
+ MISTRAL_API_KEY Model: mistral-small-latest
+ OLLAMA_HOST Ollama daemon URL (e.g. http://localhost:11434)
+ + OLLAMA_MODEL (default: llava:latest)
+
+All tests are marked ``@pytest.mark.integration``.
+
+Design
+------
+The AI provider is called **exactly once** (on the LM317 datasheet) and the
+result is cached in a module-scoped fixture. All individual tests then assert
+on this single cached result. This minimises API calls, cost, and the risk of
+hitting rate limits.
+
+Relationship to short_description / summary / abstract
+-------------------------------------------------------
+``short_description`` is FORetrieval's equivalent of a document summary or
+abstract. It is always an empty string when no AI provider is configured
+(tested in test_metadata_no_ai.py), and is filled by the LLM with 1–3
+factual sentences when AI is enabled. The test
+``test_ai_short_description_populated`` validates this behaviour.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from foretrieval.metadata import ai_metadata_provider_factory
+from foretrieval.models_metadata import DocMetadata
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+
+TESTS_DIR = Path(__file__).parent
+DATA_DIR = TESTS_DIR / "data"
+
+LM317_PDF = DATA_DIR / "lm317_voltage_regulator.pdf"
+ULN2003_PDF = DATA_DIR / "uln2003_driver.pdf"
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def ai_provider(ai_cfg):
+ """Module-scoped AI metadata provider built from the detected backend."""
+ return ai_metadata_provider_factory(ai_cfg)
+
+
+@pytest.fixture(scope="module")
+def ai_metadata_lm317(ai_provider):
+ """Single real API call for the LM317 datasheet, cached for the whole module.
+
+ All per-field tests use this fixture so we only pay for one LLM call
+ regardless of how many tests run.
+ """
+ try:
+ return ai_provider(LM317_PDF)
+ except Exception as exc:
+ # Surfaced as xfail rather than error so the rest of the suite continues
+ if _is_rate_limit(exc):
+ pytest.xfail(f"Rate limit hit during AI metadata generation: {exc}")
+ raise
+
+
+def _is_rate_limit(exc: Exception) -> bool:
+ """Heuristic: treat common rate-limit / quota errors as xfail."""
+ msg = str(exc).lower()
+ return any(
+ kw in msg for kw in ("rate limit", "429", "quota", "too many requests")
+ )
+
+
+# ---------------------------------------------------------------------------
+# Provider construction
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+def test_ai_provider_is_callable(ai_cfg):
+ """ai_metadata_provider_factory returns a callable given a valid ai_cfg."""
+ provider = ai_metadata_provider_factory(ai_cfg)
+ assert callable(provider)
+
+
+# ---------------------------------------------------------------------------
+# Base (non-AI) fields are preserved in the AI result
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+def test_ai_base_fields_preserved(ai_metadata_lm317):
+ """Non-AI filesystem fields are still present in the AI-enriched result."""
+ raw = ai_metadata_lm317
+ assert isinstance(raw.get("source_path"), str) and raw["source_path"]
+ assert raw.get("stem") == "lm317_voltage_regulator"
+ assert raw.get("ext") == ".pdf"
+ assert isinstance(raw.get("mime"), str) and raw["mime"]
+ assert isinstance(raw.get("mtime"), str) and raw["mtime"]
+
+
+@pytest.mark.integration
+def test_ai_pdf_page_count_preserved(ai_metadata_lm317):
+ """page_count extracted from the PDF is not overwritten by AI."""
+ raw = ai_metadata_lm317
+ assert isinstance(raw.get("page_count"), int)
+ assert raw["page_count"] > 0
+
+
+# ---------------------------------------------------------------------------
+# AI-enriched fields
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+def test_ai_language_populated(ai_metadata_lm317):
+ """language is a non-empty string after AI enrichment."""
+ raw = ai_metadata_lm317
+ lang = raw.get("language")
+ assert isinstance(lang, str) and len(lang) > 0, (
+ f"Expected non-empty language string, got {lang!r}"
+ )
+
+
+@pytest.mark.integration
+def test_ai_language_is_english(ai_metadata_lm317):
+ """LM317 datasheet language should be detected as English ('en').
+
+ Marked xfail in case a different LLM returns a variant code ('en-US',
+ 'english', etc.) rather than raising a hard failure.
+ """
+ raw = ai_metadata_lm317
+ lang = (raw.get("language") or "").lower().strip()
+ if not lang.startswith("en"):
+ pytest.xfail(
+ f"Language was detected as '{lang}' instead of 'en'. "
+ "This may be acceptable depending on LLM output format."
+ )
+
+
+@pytest.mark.integration
+def test_ai_tags_non_empty(ai_metadata_lm317):
+ """tags is a non-empty list after AI enrichment."""
+ raw = ai_metadata_lm317
+ tags = raw.get("tags")
+ assert isinstance(tags, list) and len(tags) > 0, (
+ f"Expected non-empty tags list, got {tags!r}"
+ )
+
+
+@pytest.mark.integration
+def test_ai_tags_are_strings(ai_metadata_lm317):
+ """Every element of tags is a string."""
+ raw = ai_metadata_lm317
+ for tag in raw.get("tags", []):
+ assert isinstance(tag, str), f"Non-string tag: {tag!r}"
+
+
+@pytest.mark.integration
+def test_ai_tags_normalized_lowercase(ai_metadata_lm317):
+ """DocMetadata normalises tags to lowercase.
+
+ The DocMetadata field_validator lowercases all tags, so even if the LLM
+ returns mixed-case tags they are stored correctly.
+ """
+ md = DocMetadata(**ai_metadata_lm317)
+ for tag in md.tags:
+ assert tag == tag.lower(), f"Tag not lowercased: {tag!r}"
+
+
+@pytest.mark.integration
+def test_ai_document_type_not_unknown(ai_metadata_lm317):
+ """document_type is classified by the AI (not the default 'unknown')."""
+ raw = ai_metadata_lm317
+ doc_type = raw.get("document_type", "")
+ assert isinstance(doc_type, str) and doc_type, "document_type is empty"
+ assert doc_type.lower() != "unknown", (
+ f"AI left document_type as 'unknown': {doc_type!r}"
+ )
+
+
+@pytest.mark.integration
+def test_ai_short_description_populated(ai_metadata_lm317):
+ """short_description (the summary/abstract equivalent) is filled by the AI.
+
+ This is the key test that validates AI-backed summary generation.
+ Without AI, short_description is always '' (tested in
+ test_metadata_no_ai.py::test_no_ai_short_description_empty).
+ With AI, the LLM produces 1-3 factual sentences — at least 20 characters.
+ """
+ raw = ai_metadata_lm317
+ desc = raw.get("short_description", "")
+ assert isinstance(desc, str), f"short_description is not a string: {desc!r}"
+ assert len(desc) >= 20, (
+ f"short_description is too short (< 20 chars): {desc!r}"
+ )
+
+
+# ---------------------------------------------------------------------------
+# DocMetadata construction
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+def test_ai_result_valid_doc_metadata(ai_metadata_lm317):
+ """DocMetadata(**raw) succeeds for the AI-enriched result."""
+ md = DocMetadata(**ai_metadata_lm317)
+ assert md.stem == "lm317_voltage_regulator"
+ assert md.ext == ".pdf"
+ assert md.language is not None
+ assert len(md.tags) > 0
+ assert md.short_description # non-empty
+
+
+@pytest.mark.integration
+def test_ai_no_hallucinated_author(ai_metadata_lm317):
+ """author is None or a plausible non-empty string.
+
+ The LM317 datasheet has no embedded PDF author metadata, so the AI
+ should either return None or a meaningful string — not a hallucinated
+ generic placeholder.
+ """
+ raw = ai_metadata_lm317
+ author = raw.get("author")
+ if author is not None:
+ assert isinstance(author, str) and len(author.strip()) > 0, (
+ f"author is set but empty/whitespace: {author!r}"
+ )
+ # Sanity: should not be an obvious placeholder
+ suspicious = {"unknown", "n/a", "none", "anonymous"}
+ assert author.strip().lower() not in suspicious, (
+ f"author looks like a hallucinated placeholder: {author!r}"
+ )
+
+
+# ---------------------------------------------------------------------------
+# Rate-limit resilience (second provider call)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.integration
+def test_ai_second_document_or_xfail(ai_provider):
+ """Run the AI provider on ULN2003 datasheet; xfail on rate limit.
+
+ This validates that the provider works for a second document without
+ exhausting the session-level cache. Rate-limit errors are tolerated
+ (xfail) since this is an additional call beyond the module fixture.
+ """
+ try:
+ raw = ai_provider(ULN2003_PDF)
+ except Exception as exc:
+ if _is_rate_limit(exc):
+ pytest.xfail(f"Rate limit hit on second document: {exc}")
+ raise
+
+ # Basic sanity on the second result
+ assert raw.get("ext") == ".pdf"
+ assert isinstance(raw.get("short_description"), str)
+ DocMetadata(**raw) # must be constructible
diff --git a/tests/test_metadata_filter.py b/tests/test_metadata_filter.py
new file mode 100644
index 0000000..b75fa6f
--- /dev/null
+++ b/tests/test_metadata_filter.py
@@ -0,0 +1,306 @@
+"""
+Tests for MetadataFilter regex support and empty-filter crash fix.
+
+No GPU, no API key, no network required — all ColPali internals are mocked.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from foretrieval.models_metadata import MetadataFilter
+from foretrieval.utils import _value_match
+
+
+# ---------------------------------------------------------------------------
+# Fixtures — representative metadata dicts
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def meta_pdf():
+ """Metadata for a PDF named 'general_datasheet.pdf'."""
+ return {
+ "stem": "general_datasheet",
+ "ext": ".pdf",
+ "title": "General Motor Controller Reference Manual",
+ "author": "Smith, J. and Doe, A.",
+ "mime": "application/pdf",
+ "language": None,
+ "tags": [],
+ "document_type": "unknown",
+ "short_description": "",
+ "mtime": "2025-03-01T10:00:00+00:00",
+ "page_count": 42,
+ }
+
+
+@pytest.fixture
+def meta_docx():
+ """Metadata for a Word document named 'safety_manual.docx'."""
+ return {
+ "stem": "safety_manual",
+ "ext": ".docx",
+ "title": "Industrial Safety Procedures",
+ "author": "Johnson, R.",
+ "mime": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+ "language": None,
+ "tags": [],
+ "document_type": "unknown",
+ "short_description": "",
+ "mtime": "2024-11-15T08:30:00+00:00",
+ "page_count": None,
+ }
+
+
+# ---------------------------------------------------------------------------
+# _value_match — regex tests
+# ---------------------------------------------------------------------------
+
+
+class TestRegexMatching:
+ def test_stem_substring_match(self, meta_pdf):
+ """regex on stem matches when the pattern is a substring."""
+ f = MetadataFilter(regex={"stem": "general"})
+ assert _value_match(meta_pdf, f) is True
+
+ def test_stem_no_match(self, meta_pdf):
+ """regex on stem returns False when pattern is not found."""
+ f = MetadataFilter(regex={"stem": "missing"})
+ assert _value_match(meta_pdf, f) is False
+
+ def test_title_case_insensitive(self, meta_pdf):
+ """regex matching is case-insensitive by default."""
+ f = MetadataFilter(regex={"title": "GENERAL"})
+ assert _value_match(meta_pdf, f) is True
+
+ def test_author_substring(self, meta_pdf):
+ """regex on author matches a substring of the author string."""
+ f = MetadataFilter(regex={"author": "smith"})
+ assert _value_match(meta_pdf, f) is True
+
+ def test_author_no_match(self, meta_pdf):
+ """regex on author returns False when pattern is not in author string."""
+ f = MetadataFilter(regex={"author": "nobody"})
+ assert _value_match(meta_pdf, f) is False
+
+ def test_alternation(self, meta_pdf):
+ """regex alternation (|) matches if any alternative is present."""
+ f = MetadataFilter(regex={"title": "motor|pump"})
+ assert _value_match(meta_pdf, f) is True
+
+ def test_alternation_no_match(self, meta_pdf):
+ """regex alternation returns False when none of the alternatives match."""
+ f = MetadataFilter(regex={"title": "bicycle|aircraft"})
+ assert _value_match(meta_pdf, f) is False
+
+ def test_anchored_start_match(self, meta_pdf):
+ """regex ^ anchor matches at the start of the field value."""
+ f = MetadataFilter(regex={"stem": "^general"})
+ assert _value_match(meta_pdf, f) is True
+
+ def test_anchored_start_no_match(self, meta_docx):
+ """regex ^ anchor fails when value does not start with the pattern."""
+ f = MetadataFilter(regex={"stem": "^general"})
+ assert _value_match(meta_docx, f) is False
+
+ def test_malformed_pattern_no_crash(self, meta_pdf):
+ """A malformed regex pattern returns False without raising."""
+ f = MetadataFilter(regex={"stem": "[invalid"})
+ result = _value_match(meta_pdf, f)
+ assert result is False
+
+ def test_field_not_in_meta_no_match(self, meta_pdf):
+ """regex on a field absent from metadata returns False."""
+ f = MetadataFilter(regex={"nonexistent_field": "anything"})
+ assert _value_match(meta_pdf, f) is False
+
+ def test_regex_on_ai_field_language(self, meta_pdf):
+ """regex works on AI fields like language when present."""
+ meta = dict(meta_pdf)
+ meta["language"] = "en"
+ f = MetadataFilter(regex={"language": "^en"})
+ assert _value_match(meta, f) is True
+
+ def test_regex_on_document_type(self, meta_pdf):
+ """regex works on document_type."""
+ meta = dict(meta_pdf)
+ meta["document_type"] = "technical note"
+ f = MetadataFilter(regex={"document_type": "technical"})
+ assert _value_match(meta, f) is True
+
+ def test_regex_combined_with_ext_and(self, meta_pdf, meta_docx):
+ """regex AND ext: only matches when both conditions hold."""
+ f = MetadataFilter(ext=".pdf", regex={"stem": "general"})
+ assert _value_match(meta_pdf, f) is True # PDF + stem contains "general"
+ assert _value_match(meta_docx, f) is False # docx fails ext check
+
+ def test_regex_combined_with_ext_or(self, meta_pdf, meta_docx):
+ """regex OR ext: matches when either condition holds."""
+ f = MetadataFilter(ext=".docx", regex={"stem": "general"}, logic="OR")
+ assert _value_match(meta_pdf, f) is True # stem matches
+ assert _value_match(meta_docx, f) is True # ext matches
+
+ def test_regex_not_double_processed_by_extra_loop(self, meta_pdf):
+ """The 'regex' key must not be processed by the extra-field equality loop.
+
+ If it were, 'regex' as a key would be compared against meta['regex']
+ which doesn't exist, producing a spurious False check.
+ """
+ f = MetadataFilter(regex={"stem": "general"})
+ # Ensure no check is generated for a literal 'regex' key equality
+ # The filter has one regex condition → exactly one check in the list
+ # We verify this by asserting the match is True (would be False if
+ # the extra loop produced a False check for missing 'regex' field).
+ assert _value_match(meta_pdf, f) is True
+
+ def test_multiple_regex_fields_and(self, meta_pdf):
+ """Multiple regex entries are all required under AND logic."""
+ f = MetadataFilter(regex={"stem": "general", "title": "motor"})
+ assert _value_match(meta_pdf, f) is True
+
+ def test_multiple_regex_fields_one_fails(self, meta_pdf):
+ """Multiple regex entries: one failing under AND logic → False."""
+ f = MetadataFilter(regex={"stem": "general", "title": "aircraft"})
+ assert _value_match(meta_pdf, f) is False
+
+ def test_multiple_regex_fields_or(self, meta_pdf):
+ """Multiple regex entries under OR: one match is sufficient."""
+ f = MetadataFilter(regex={"stem": "general", "title": "aircraft"}, logic="OR")
+ assert _value_match(meta_pdf, f) is True
+
+
+# ---------------------------------------------------------------------------
+# Empty filter → [] and no crash
+# ---------------------------------------------------------------------------
+
+
+class TestEmptyFilterReturnsEmptyList:
+ """Tests that ColPaliModel.search() returns [] when the filter matches nothing."""
+
+ def _make_mock_colpali(self):
+ """Return a minimal mock ColPaliModel with in-memory index state."""
+ from foretrieval.colpali import ColPaliModel
+
+ model = MagicMock(spec=ColPaliModel)
+ model.device = "cpu"
+ model.verbose = 0
+ model.collection = None
+ model.enable_heatmaps = False
+ model.enable_circle = False
+ model.storage_qdrant = False # use local search path
+ model._encode_search_query = MagicMock(return_value=[MagicMock()])
+
+ # One document indexed, stored with ext=".pdf"
+ model.doc_id_to_metadata = {
+ 0: {"stem": "datasheet", "ext": ".pdf", "title": "Test Doc"}
+ }
+ model.embed_id_to_doc_id = {0: {"doc_id": 0, "page_id": 0}}
+ model.indexed_embeddings = [MagicMock()] # one page embedding
+
+ # processor mock
+ processor = MagicMock()
+ processor.process_queries.return_value = {
+ "input_ids": MagicMock(
+ __getitem__=lambda self, idx: MagicMock(
+ detach=lambda: MagicMock(
+ cpu=lambda: MagicMock(tolist=lambda: [1, 2, 3])
+ )
+ )
+ )
+ }
+ processor.tokenizer.convert_ids_to_tokens.return_value = ["tok1", "tok2", "tok3"]
+ model.processor = processor
+
+ # model forward pass mock
+ inner_model = MagicMock()
+ inner_model.dtype = None
+ model.model = inner_model
+
+ return model
+
+ def test_filter_no_match_returns_empty_list(self):
+ """LocalVectorStore.search() returns [] when the metadata filter matches nothing."""
+ from foretrieval.vector_store.local import LocalVectorStore
+ from foretrieval.vector_store.base import MultiVectorQuery, StoredPoint, make_point_id
+ import torch
+ from unittest.mock import MagicMock
+
+ store = LocalVectorStore()
+
+ import tempfile, pathlib
+ with tempfile.TemporaryDirectory() as tmp:
+ store.open("idx", pathlib.Path(tmp), create=True)
+ store.set_doc_id_to_metadata({0: {"stem": "datasheet", "ext": ".pdf"}})
+ # Insert one embedding
+ store.upsert([StoredPoint(
+ point_id=make_point_id(0, 1),
+ vector=torch.rand(4, 8),
+ payload={"doc_id": 0, "page_id": 1, "chunk_id": None, "metadata": {}},
+ )])
+ proc = MagicMock()
+ store.set_processor(proc)
+
+ # Filter on language=fr which doesn't match → empty result
+ q = MultiVectorQuery(vectors=torch.rand(2, 8), filter_metadata={"language": "fr"})
+ results = store.search(q, k=5)
+ assert results == []
+
+ def test_filter_no_match_does_not_raise(self):
+ """search() must not raise ValueError when filter matches nothing."""
+ from foretrieval.vector_store.local import LocalVectorStore
+ from foretrieval.vector_store.base import MultiVectorQuery, StoredPoint, make_point_id
+ import torch
+ from unittest.mock import MagicMock
+
+ store = LocalVectorStore()
+ import tempfile, pathlib
+ with tempfile.TemporaryDirectory() as tmp:
+ store.open("idx", pathlib.Path(tmp), create=True)
+ store.set_doc_id_to_metadata({0: {"ext": ".pdf"}})
+ store.upsert([StoredPoint(
+ point_id=make_point_id(0, 1),
+ vector=torch.rand(4, 8),
+ payload={"doc_id": 0, "page_id": 1, "chunk_id": None, "metadata": {}},
+ )])
+ proc = MagicMock()
+ store.set_processor(proc)
+ q = MultiVectorQuery(vectors=torch.rand(2, 8), filter_metadata={"language": "fr"})
+ try:
+ store.search(q, k=5)
+ except ValueError as exc:
+ pytest.fail(f"search raised ValueError when filter matched nothing: {exc}")
+
+ def test_filter_match_still_works(self):
+ """search() returns results when the filter matches documents."""
+ from foretrieval.colpali import ColPaliModel
+ from unittest.mock import MagicMock
+ from foretrieval.vector_store.base import SearchHit
+ from foretrieval.objects import Result
+
+ mock = MagicMock()
+ mock.collection = {}
+ mock.enable_heatmaps = False
+ mock.enable_circle = False
+ mock.doc_id_to_metadata = {0: {"ext": ".pdf"}}
+ mock._encode_search_query = MagicMock(return_value=[MagicMock()])
+
+ hit = SearchHit(point_id=0, score=0.9, payload={"doc_id": 0, "page_id": 1, "chunk_id": None, "metadata": {}})
+ mock.vector_store.search.return_value = [hit]
+ mock.embed_id_to_extra = {}
+ # Make _hits_to_results and _finalize_results return proper values
+ expected = [Result(doc_id=0, page_num=1, score=0.9)]
+ mock._hits_to_results = MagicMock(return_value=expected)
+ mock._finalize_results = MagicMock(return_value=expected)
+
+ results = ColPaliModel.search(
+ mock,
+ query="test",
+ k=1,
+ filter_metadata={"ext": ".pdf"},
+ return_base64_results=False,
+ )
+ assert isinstance(results, list)
+ assert len(results) > 0
diff --git a/tests/test_metadata_no_ai.py b/tests/test_metadata_no_ai.py
new file mode 100644
index 0000000..46c3a31
--- /dev/null
+++ b/tests/test_metadata_no_ai.py
@@ -0,0 +1,378 @@
+"""
+Tests for FORetrieval metadata generation WITHOUT an AI provider.
+
+These tests exercise the ``_no_ai_provider`` callable returned by
+``ai_metadata_provider_factory(None)``. No API key, no GPU, and no network
+access are required — every test runs against real files on disk.
+
+Metadata fields and their expected values without AI
+-----------------------------------------------------
++--------------------+----------------------------------+
+| Field | Expected value (no AI) |
++====================+==================================+
+| source_path | str (absolute path) |
+| stem | str (filename without extension) |
+| ext | str (starts with '.', lowercase) |
+| mime | str |
+| mtime | ISO-8601 UTC string |
+| page_count | int > 0 (PDFs only) |
+| image_width | int > 0 (images only) |
+| image_height | int > 0 (images only) |
+| author | str or None (from PDF metadata) |
+| title | str or None (from PDF metadata) |
+| language | None ← not detected without AI |
+| tags | [] ← not generated without AI |
+| document_type | "unknown" |
+| short_description | "" ← the "summary/abstract" |
+| | equivalent; empty without AI |
++--------------------+----------------------------------+
+
+Note: ``short_description`` is the field that corresponds to a document
+summary or abstract. It is empty when no AI provider is configured, and
+is filled by the LLM when an AI provider is used (see test_metadata_ai.py).
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from pathlib import Path
+
+import pytest
+
+from foretrieval.metadata import ai_metadata_provider_factory
+from foretrieval.models_metadata import DocMetadata, build_metadata_list_for_dir
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+
+TESTS_DIR = Path(__file__).parent
+DATA_DIR = TESTS_DIR / "data"
+SAMPLE_DATA_DIR = TESTS_DIR.parent / "sample_data"
+
+LM317_PDF = DATA_DIR / "lm317_voltage_regulator.pdf"
+ULN2003_PDF = DATA_DIR / "uln2003_driver.pdf"
+SAMPLE_PNG = SAMPLE_DATA_DIR / "sample_png.png"
+SAMPLE_JPG = SAMPLE_DATA_DIR / "sample_jpg.jpg"
+SAMPLE_TXT = SAMPLE_DATA_DIR / "sample_txt.txt"
+SAMPLE_DOCX = SAMPLE_DATA_DIR / "sample_docx.docx"
+SAMPLE_XLSX = SAMPLE_DATA_DIR / "sample_xlsx.xlsx"
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def provider():
+ """Module-scoped no-AI metadata provider."""
+ return ai_metadata_provider_factory(None)
+
+
+# ---------------------------------------------------------------------------
+# Base field tests
+# ---------------------------------------------------------------------------
+
+
+def test_no_ai_pdf_base_fields(provider):
+ """All base filesystem fields are populated for a PDF."""
+ raw = provider(LM317_PDF)
+
+ assert isinstance(raw["source_path"], str)
+ assert raw["source_path"].endswith("lm317_voltage_regulator.pdf")
+ assert raw["stem"] == "lm317_voltage_regulator"
+ assert isinstance(raw["mime"], str) and raw["mime"]
+ assert isinstance(raw["mtime"], str) and raw["mtime"]
+
+
+def test_no_ai_ext_normalized_pdf(provider):
+ """ext is '.pdf' (lowercase, starts with dot) for a PDF file."""
+ raw = provider(LM317_PDF)
+ assert raw["ext"] == ".pdf"
+
+
+def test_no_ai_ext_normalized_image(provider):
+ """ext starts with '.' and is lowercase for image files."""
+ for path in [SAMPLE_PNG, SAMPLE_JPG]:
+ raw = provider(path)
+ assert raw["ext"].startswith("."), f"ext missing leading dot for {path.name}"
+ assert raw["ext"] == raw["ext"].lower(), f"ext not lowercase for {path.name}"
+
+
+def test_no_ai_mtime_is_valid_iso_utc(provider):
+ """mtime is a parseable ISO-8601 string with UTC timezone info."""
+ raw = provider(LM317_PDF)
+ mtime_str = raw["mtime"]
+ # Parse via datetime — should not raise
+ dt = datetime.fromisoformat(mtime_str.replace("Z", "+00:00"))
+ assert dt.tzinfo is not None, "mtime must carry timezone info"
+
+
+# ---------------------------------------------------------------------------
+# PDF-specific tests
+# ---------------------------------------------------------------------------
+
+
+def test_no_ai_pdf_page_count_lm317(provider):
+ """page_count is a positive integer for the LM317 datasheet."""
+ raw = provider(LM317_PDF)
+ assert isinstance(raw["page_count"], int)
+ assert raw["page_count"] > 0
+
+
+def test_no_ai_pdf_page_count_uln2003(provider):
+ """page_count is a positive integer for the ULN2003 datasheet."""
+ raw = provider(ULN2003_PDF)
+ assert isinstance(raw["page_count"], int)
+ assert raw["page_count"] > 0
+
+
+def test_no_ai_non_pdf_page_count_is_none(provider):
+ """page_count is None for non-PDF files."""
+ for path in [SAMPLE_PNG, SAMPLE_TXT, SAMPLE_DOCX]:
+ raw = provider(path)
+ assert raw["page_count"] is None, (
+ f"Expected page_count=None for {path.name}, got {raw['page_count']}"
+ )
+
+
+# ---------------------------------------------------------------------------
+# Image dimension tests
+# ---------------------------------------------------------------------------
+
+
+def test_no_ai_image_dims_png(provider):
+ """image_width and image_height are positive ints for a PNG."""
+ raw = provider(SAMPLE_PNG)
+ assert isinstance(raw["image_width"], int) and raw["image_width"] > 0
+ assert isinstance(raw["image_height"], int) and raw["image_height"] > 0
+
+
+def test_no_ai_image_dims_jpg(provider):
+ """image_width and image_height are positive ints for a JPEG."""
+ raw = provider(SAMPLE_JPG)
+ assert isinstance(raw["image_width"], int) and raw["image_width"] > 0
+ assert isinstance(raw["image_height"], int) and raw["image_height"] > 0
+
+
+def test_no_ai_non_image_dims_are_none(provider):
+ """image_width and image_height are None for non-image files."""
+ for path in [LM317_PDF, SAMPLE_TXT, SAMPLE_DOCX]:
+ raw = provider(path)
+ assert raw["image_width"] is None, f"Expected None for {path.name}"
+ assert raw["image_height"] is None, f"Expected None for {path.name}"
+
+
+# ---------------------------------------------------------------------------
+# AI field absence tests (the critical no-AI assertions)
+# ---------------------------------------------------------------------------
+
+
+def test_no_ai_language_is_none(provider):
+ """language is None without AI — language detection requires the LLM."""
+ for path in [LM317_PDF, ULN2003_PDF]:
+ raw = provider(path)
+ assert raw["language"] is None, (
+ f"Expected language=None without AI for {path.name}"
+ )
+
+
+def test_no_ai_tags_empty(provider):
+ """tags is an empty list without AI — tagging requires the LLM."""
+ for path in [LM317_PDF, ULN2003_PDF]:
+ raw = provider(path)
+ assert raw["tags"] == [], (
+ f"Expected tags=[] without AI for {path.name}"
+ )
+
+
+def test_no_ai_document_type_unknown(provider):
+ """document_type is 'unknown' without AI — classification requires the LLM."""
+ for path in [LM317_PDF, ULN2003_PDF]:
+ raw = provider(path)
+ assert raw["document_type"] == "unknown", (
+ f"Expected document_type='unknown' without AI for {path.name}"
+ )
+
+
+def test_no_ai_short_description_empty(provider):
+ """short_description (the summary/abstract equivalent) is empty without AI.
+
+ ``short_description`` is FORetrieval's equivalent of a document
+ summary or abstract. Without an AI provider it is always an empty
+ string; the LLM fills it with 1-3 factual sentences when AI is
+ enabled (see test_metadata_ai.py::test_ai_short_description_populated).
+ """
+ for path in [LM317_PDF, ULN2003_PDF]:
+ raw = provider(path)
+ assert raw["short_description"] == "", (
+ f"Expected short_description='' without AI for {path.name}"
+ )
+
+
+# ---------------------------------------------------------------------------
+# DocMetadata construction tests
+# ---------------------------------------------------------------------------
+
+
+def test_no_ai_result_valid_doc_metadata_lm317(provider):
+ """DocMetadata(**raw) succeeds for the LM317 datasheet."""
+ raw = provider(LM317_PDF)
+ md = DocMetadata(**raw)
+ assert md.stem == "lm317_voltage_regulator"
+ assert md.ext == ".pdf"
+
+
+def test_no_ai_result_valid_doc_metadata_uln2003(provider):
+ """DocMetadata(**raw) succeeds for the ULN2003 datasheet."""
+ raw = provider(ULN2003_PDF)
+ md = DocMetadata(**raw)
+ assert md.ext == ".pdf"
+
+
+def test_no_ai_doc_metadata_tags_normalized(provider):
+ """DocMetadata normalises tags to lowercase stripped strings."""
+ raw = provider(LM317_PDF)
+ raw["tags"] = [" Python ", "RAG", "FOO"]
+ md = DocMetadata(**raw)
+ assert md.tags == ["python", "rag", "foo"]
+
+
+def test_no_ai_doc_metadata_ext_normalized(provider):
+ """DocMetadata normalises ext: adds leading dot, lowercases."""
+ raw = provider(LM317_PDF)
+ raw["ext"] = "PDF" # deliberate unnormalized value
+ md = DocMetadata(**raw)
+ assert md.ext == ".pdf"
+
+
+# ---------------------------------------------------------------------------
+# Multiple format smoke tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "path",
+ [
+ SAMPLE_TXT,
+ SAMPLE_DOCX,
+ SAMPLE_XLSX,
+ SAMPLE_PNG,
+ SAMPLE_JPG,
+ ],
+ ids=["txt", "docx", "xlsx", "png", "jpg"],
+)
+def test_no_ai_multiple_formats_no_crash(provider, path):
+ """provider() returns a valid dict for all common file types without crashing."""
+ raw = provider(path)
+ assert isinstance(raw, dict)
+ # Must always contain these base keys
+ for key in ("source_path", "stem", "ext", "mime", "mtime"):
+ assert key in raw, f"Missing key '{key}' for {path.name}"
+ # Must be constructible as DocMetadata
+ DocMetadata(**raw)
+
+
+# ---------------------------------------------------------------------------
+# build_metadata_list_for_dir tests
+# ---------------------------------------------------------------------------
+
+
+def test_build_metadata_list_for_dir_length(provider):
+ """Result list has the same length as the recursive file enumeration used by index()."""
+ items = sorted(
+ (p for p in DATA_DIR.rglob("*") if p.is_file()),
+ key=lambda p: p.relative_to(DATA_DIR),
+ )
+ md_list = build_metadata_list_for_dir(DATA_DIR, provider)
+ assert len(md_list) == len(items)
+
+
+def test_build_metadata_list_for_dir_all_doc_metadata(provider):
+ """Every entry in the result list is a DocMetadata instance (no None placeholders)."""
+ md_list = build_metadata_list_for_dir(DATA_DIR, provider)
+ for md in md_list:
+ # All entries are files (directories are excluded); provider always
+ # returns metadata for PDF files in DATA_DIR.
+ assert isinstance(md, DocMetadata), (
+ f"Expected DocMetadata for all entries, got {type(md)}"
+ )
+
+
+def test_build_metadata_list_for_dir_order_stable(provider):
+ """Two consecutive calls return stems in the same order (sort is deterministic)."""
+ stems_first = [
+ md.stem for md in build_metadata_list_for_dir(DATA_DIR, provider)
+ if md is not None
+ ]
+ stems_second = [
+ md.stem for md in build_metadata_list_for_dir(DATA_DIR, provider)
+ if md is not None
+ ]
+ assert stems_first == stems_second, (
+ "build_metadata_list_for_dir returned different ordering on consecutive calls"
+ )
+
+
+def test_build_metadata_list_for_dir_stems_match_sorted_filenames(provider):
+ """Stems in the result list match the stems of recursively-sorted filenames.
+
+ Validates that the list is aligned with ColPaliModel.index() which iterates
+ files via rglob("*") sorted by relative path.
+ """
+ sorted_files = sorted(
+ (p for p in DATA_DIR.rglob("*") if p.is_file()),
+ key=lambda p: p.relative_to(DATA_DIR),
+ )
+ md_list = [
+ md for md in build_metadata_list_for_dir(DATA_DIR, provider)
+ if md is not None
+ ]
+ assert len(md_list) == len(sorted_files)
+ for expected_path, md in zip(sorted_files, md_list):
+ assert md.stem == expected_path.stem, (
+ f"Stem mismatch: expected '{expected_path.stem}', got '{md.stem}'"
+ )
+
+
+def test_build_metadata_list_for_dir_recursive_alignment(provider, tmp_path):
+ """metadata list is aligned with index()'s rglob enumeration for nested dirs.
+
+ This is a regression test for the bug where build_metadata_list_for_dir used
+ non-recursive iterdir() while index() used recursive rglob(), causing a count
+ mismatch and ValueError when add_metadata=True on a multi-level dataset.
+ """
+ # Create a nested directory structure:
+ # tmp_path/
+ # a.pdf (top-level file)
+ # sub/
+ # b.pdf
+ # deep/
+ # c.pdf
+ (tmp_path / "a.pdf").write_bytes(b"%PDF-1.4 fake")
+ (tmp_path / "sub").mkdir()
+ (tmp_path / "sub" / "b.pdf").write_bytes(b"%PDF-1.4 fake")
+ (tmp_path / "sub" / "deep").mkdir()
+ (tmp_path / "sub" / "deep" / "c.pdf").write_bytes(b"%PDF-1.4 fake")
+
+ def simple_provider(path):
+ return {"stem": path.stem, "ext": path.suffix}
+
+ md_list = build_metadata_list_for_dir(tmp_path, simple_provider)
+
+ # index() would enumerate: a.pdf, sub/b.pdf, sub/deep/c.pdf (sorted by relative path)
+ expected_files = sorted(
+ (p for p in tmp_path.rglob("*") if p.is_file()),
+ key=lambda p: p.relative_to(tmp_path),
+ )
+ assert len(md_list) == len(expected_files) == 3, (
+ f"Expected 3 entries, got {len(md_list)}"
+ )
+ for i, (expected_path, md) in enumerate(zip(expected_files, md_list)):
+ assert md is not None, f"Entry {i} ({expected_path.name}) should not be None"
+ assert md.stem == expected_path.stem, (
+ f"Entry {i}: stem mismatch — expected '{expected_path.stem}', got '{md.stem}'"
+ )
+
+
diff --git a/tests/test_qdrant.py b/tests/test_qdrant.py
new file mode 100644
index 0000000..8321049
--- /dev/null
+++ b/tests/test_qdrant.py
@@ -0,0 +1,445 @@
+"""
+Tests for the Qdrant storage backend and the docling ingestion guard.
+
+Since Qdrant logic now lives in foretrieval/vector_store/qdrant.py, the
+low-level tests (make_point_id, filter building, upsert shape, etc.) live in
+test_vector_store_qdrant.py. This file retains the tests that exercise
+ColPaliModel-level behaviour: storage_backend selection, index config
+serialization, search dispatch, and the docling guard.
+
+The slow integration test (marked @pytest.mark.slow) is preserved unchanged
+from the original test_qdrant.py.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from foretrieval.colpali import ColPaliModel, _DOCLING_AVAILABLE
+from foretrieval.vector_store import LocalVectorStore, QdrantVectorStore, make_point_id
+from foretrieval.vector_store.qdrant import _QDRANT_AVAILABLE
+
+DATA_DIR = Path(__file__).parent / "data"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_mock_model(storage_backend: str = "qdrant") -> ColPaliModel:
+ """Build a ColPaliModel with all heavy attributes mocked away."""
+ from foretrieval.vector_store import make_vector_store
+ model = MagicMock()
+ model.storage_backend = storage_backend
+ model.storage_qdrant = (storage_backend == "qdrant")
+ model.index_name = "test_index"
+ model.index_root = ".foretrieval_test"
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.collection = {}
+ model.embed_id_to_extra = {}
+ model.device = "cpu"
+ model.verbose = 0
+ model.processor = MagicMock()
+ model.full_document_collection = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.model_name = "stub"
+ model.vector_store = make_vector_store(storage_backend)
+ return model
+
+
+# ---------------------------------------------------------------------------
+# Docling ingestion guard
+# ---------------------------------------------------------------------------
+
+class TestDoclingGuard:
+ def test_docling_not_available_raises_on_init(self):
+ """When docling is absent, requesting docling backend raises RuntimeError."""
+ with (
+ patch("foretrieval.colpali._DOCLING_AVAILABLE", False),
+ patch("foretrieval.colpali.ColPaliModel._load_model_and_processor"),
+ patch("foretrieval.colpali.ColPaliModel._load_index_state"),
+ ):
+ with pytest.raises(RuntimeError, match="docling"):
+ ColPaliModel(
+ pretrained_model_name_or_path="colqwen2-stub",
+ ingestion={"backend": "docling"},
+ index_root="/tmp",
+ )
+
+ def test_docling_not_available_message_contains_install_hint(self):
+ """RuntimeError for missing docling contains pip install hint."""
+ with (
+ patch("foretrieval.colpali._DOCLING_AVAILABLE", False),
+ patch("foretrieval.colpali.ColPaliModel._load_model_and_processor"),
+ patch("foretrieval.colpali.ColPaliModel._load_index_state"),
+ ):
+ with pytest.raises(RuntimeError) as exc_info:
+ ColPaliModel(
+ pretrained_model_name_or_path="colqwen2-stub",
+ ingestion={"backend": "docling"},
+ index_root="/tmp",
+ )
+ assert "foretrieval[docling]" in str(exc_info.value)
+
+
+# ---------------------------------------------------------------------------
+# make_point_id (now in vector_store.base, still accessible publicly)
+# ---------------------------------------------------------------------------
+
+class TestMakePointId:
+ def test_basic(self):
+ result = make_point_id(doc_id=1, page_id=2, chunk_id=3)
+ expected = 1 * 10_000_000 + 2 * 10_000 + 3
+ assert result == expected
+
+ def test_none_chunk_uses_zero(self):
+ result_none = make_point_id(doc_id=1, page_id=2, chunk_id=None)
+ result_zero = make_point_id(doc_id=1, page_id=2, chunk_id=0)
+ assert result_none == result_zero
+
+ def test_uniqueness(self):
+ ids = {
+ make_point_id(d, p, c)
+ for d, p, c in [(0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0)]
+ }
+ assert len(ids) == 4
+
+ def test_deterministic(self):
+ a = make_point_id(5, 3, 7)
+ b = make_point_id(5, 3, 7)
+ assert a == b
+
+
+# ---------------------------------------------------------------------------
+# QdrantVectorStore._build_filter (moved from ColPaliModel._build_qdrant_filter)
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _QDRANT_AVAILABLE, reason="qdrant-client not installed")
+class TestBuildQdrantFilter:
+ def _store(self) -> QdrantVectorStore:
+ store = QdrantVectorStore.__new__(QdrantVectorStore)
+ store._client = MagicMock()
+ store._collection_name = "test_col"
+ store._index_root = Path("/tmp")
+ return store
+
+ def test_none_returns_none(self):
+ assert self._store()._build_filter(None) is None
+
+ def test_empty_dict_returns_none(self):
+ assert self._store()._build_filter({}) is None
+
+ def test_single_field_returns_filter(self):
+ from qdrant_client.models import Filter, FieldCondition
+ result = self._store()._build_filter({"ext": ".pdf"})
+ assert isinstance(result, Filter)
+ assert len(result.must) == 1
+ cond = result.must[0]
+ assert isinstance(cond, FieldCondition)
+ assert cond.key == "metadata.ext"
+ assert cond.match.value == ".pdf"
+
+ def test_multiple_fields_returns_multiple_conditions(self):
+ from qdrant_client.models import Filter
+ result = self._store()._build_filter({"ext": ".pdf", "language": "en"})
+ assert isinstance(result, Filter)
+ assert len(result.must) == 2
+ keys = {c.key for c in result.must}
+ assert "metadata.ext" in keys
+ assert "metadata.language" in keys
+
+
+# ---------------------------------------------------------------------------
+# Search dispatch — now goes through vector_store.search()
+# ---------------------------------------------------------------------------
+
+class TestSearchDispatch:
+ def test_zero_k_returns_empty_immediately(self):
+ """search() returns [] when k < 1 without calling any search backend."""
+ mock = _make_mock_model("local")
+ mock._encode_search_query = MagicMock()
+ result = ColPaliModel.search(mock, query="test", k=0)
+ assert result == []
+ mock._encode_search_query.assert_not_called()
+
+ def test_search_delegates_to_vector_store(self):
+ """search() calls self.vector_store.search() regardless of backend."""
+ mock = _make_mock_model("local")
+ mock.vector_store = MagicMock()
+ mock.vector_store.search.return_value = []
+ mock._encode_search_query = MagicMock(return_value=[MagicMock()])
+ ColPaliModel.search(mock, query="test", k=3)
+ mock.vector_store.search.assert_called_once()
+
+ def test_search_qdrant_backend_uses_qdrant_store(self):
+ """ColPaliModel with qdrant backend delegates to QdrantVectorStore."""
+ mock = _make_mock_model("qdrant")
+ assert isinstance(mock.vector_store, QdrantVectorStore)
+
+ def test_search_local_backend_uses_local_store(self):
+ """ColPaliModel with local backend delegates to LocalVectorStore."""
+ mock = _make_mock_model("local")
+ assert isinstance(mock.vector_store, LocalVectorStore)
+
+
+# ---------------------------------------------------------------------------
+# Empty filter → empty results (through LocalVectorStore)
+# ---------------------------------------------------------------------------
+
+class TestEmptyFilterLocal:
+ def test_local_empty_filter_returns_empty_list(self, tmp_path):
+ """When LocalVectorStore metadata filter matches nothing, search returns []."""
+ from foretrieval.vector_store.base import MultiVectorQuery
+ import torch
+
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ # Empty metadata map → filter matches nothing
+ store.set_doc_id_to_metadata({})
+ proc = MagicMock()
+ proc.score.return_value = torch.tensor([[]])
+ store.set_processor(proc)
+
+ q = MultiVectorQuery(
+ vectors=torch.rand(2, 8),
+ filter_metadata={"language": "fr"},
+ )
+ result = store.search(q, k=5)
+ assert result == []
+
+ def test_local_empty_filter_does_not_raise(self, tmp_path):
+ """search() must not raise ValueError when filter matches nothing."""
+ from foretrieval.vector_store.base import MultiVectorQuery
+ import torch
+
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ store.set_doc_id_to_metadata({})
+ proc = MagicMock()
+ store.set_processor(proc)
+
+ q = MultiVectorQuery(
+ vectors=torch.rand(2, 8),
+ filter_metadata={"language": "fr"},
+ )
+ try:
+ store.search(q, k=5)
+ except ValueError as exc:
+ pytest.fail(f"search raised ValueError on empty filter: {exc}")
+
+
+# ---------------------------------------------------------------------------
+# storage_backend serialization — index_config.json.gz
+# ---------------------------------------------------------------------------
+
+class TestStorageBackendSerialization:
+ def test_qdrant_backend_written_to_config(self, tmp_path):
+ """_export_index writes storage_backend='qdrant'."""
+ import srsly
+ mock = _make_mock_model("qdrant")
+ mock.index_root = str(tmp_path)
+ mock.index_name = "test_idx"
+ mock.full_document_collection = False
+ mock.highest_doc_id = 0
+ mock.max_image_width = None
+ mock.max_image_height = None
+ mock.embed_id_to_extra = {}
+ mock.doc_ids_to_file_names = {}
+ mock.doc_id_to_metadata = {}
+ mock.collection = {}
+ mock.vector_store = MagicMock()
+
+ ColPaliModel._export_index(mock)
+
+ config = srsly.read_gzip_json(tmp_path / "test_idx" / "index_config.json.gz")
+ assert config["storage_backend"] == "qdrant"
+
+ def test_local_backend_written_to_config(self, tmp_path):
+ """_export_index writes storage_backend='local'."""
+ import srsly
+ mock = _make_mock_model("local")
+ mock.index_root = str(tmp_path)
+ mock.index_name = "test_idx"
+ mock.full_document_collection = False
+ mock.highest_doc_id = 0
+ mock.max_image_width = None
+ mock.max_image_height = None
+ mock.embed_id_to_extra = {}
+ mock.doc_ids_to_file_names = {}
+ mock.doc_id_to_metadata = {}
+ mock.collection = {}
+ mock.vector_store = MagicMock()
+
+ ColPaliModel._export_index(mock)
+
+ config = srsly.read_gzip_json(tmp_path / "test_idx" / "index_config.json.gz")
+ assert config["storage_backend"] == "local"
+
+ def test_milvus_backend_written_to_config(self, tmp_path):
+ """_export_index writes storage_backend='milvus'."""
+ import srsly
+ mock = _make_mock_model("milvus")
+ mock.index_root = str(tmp_path)
+ mock.index_name = "test_idx"
+ mock.full_document_collection = False
+ mock.highest_doc_id = 0
+ mock.max_image_width = None
+ mock.max_image_height = None
+ mock.embed_id_to_extra = {}
+ mock.doc_ids_to_file_names = {}
+ mock.doc_id_to_metadata = {}
+ mock.collection = {}
+ mock.vector_store = MagicMock()
+
+ ColPaliModel._export_index(mock)
+
+ config = srsly.read_gzip_json(tmp_path / "test_idx" / "index_config.json.gz")
+ assert config["storage_backend"] == "milvus"
+
+ def test_from_index_reads_qdrant_backend(self, tmp_path):
+ """from_index() passes storage_backend='qdrant' when config says so."""
+ import srsly, torch
+
+ idx_path = tmp_path / "my_index"
+ idx_path.mkdir()
+ srsly.write_gzip_json(
+ idx_path / "index_config.json.gz",
+ {
+ "model_name": "vidore/colqwen2.5-v0.2",
+ "storage_backend": "qdrant",
+ "full_document_collection": False,
+ "highest_doc_id": 0,
+ "resize_stored_images": False,
+ "max_image_width": None,
+ "max_image_height": None,
+ },
+ )
+ torch.save({}, idx_path / "embed_id_to_extra.pt")
+ srsly.write_gzip_json(idx_path / "doc_ids_to_file_names.json.gz", {})
+ srsly.write_gzip_json(idx_path / "metadata.json.gz", {})
+
+ with patch("foretrieval.colpali.ColPaliModel.__init__", return_value=None) as mock_init:
+ try:
+ ColPaliModel.from_index(
+ index_path=str(idx_path.name),
+ index_root=str(tmp_path),
+ device="cpu",
+ )
+ except Exception:
+ pass
+
+ if mock_init.called:
+ call_kwargs = mock_init.call_args.kwargs
+ assert call_kwargs.get("storage_backend") == "qdrant"
+
+ def test_from_index_reads_local_backend(self, tmp_path):
+ """from_index() passes storage_backend='local' when config says so."""
+ import srsly, torch
+
+ idx_path = tmp_path / "local_index"
+ idx_path.mkdir()
+ srsly.write_gzip_json(
+ idx_path / "index_config.json.gz",
+ {
+ "model_name": "vidore/colqwen2.5-v0.2",
+ "storage_backend": "local",
+ "full_document_collection": False,
+ "highest_doc_id": 0,
+ "resize_stored_images": False,
+ "max_image_width": None,
+ "max_image_height": None,
+ },
+ )
+ torch.save({}, idx_path / "embed_id_to_extra.pt")
+ srsly.write_gzip_json(idx_path / "doc_ids_to_file_names.json.gz", {})
+ srsly.write_gzip_json(idx_path / "metadata.json.gz", {})
+
+ with patch("foretrieval.colpali.ColPaliModel.__init__", return_value=None) as mock_init:
+ try:
+ ColPaliModel.from_index(
+ index_path=str(idx_path.name),
+ index_root=str(tmp_path),
+ device="cpu",
+ )
+ except Exception:
+ pass
+
+ if mock_init.called:
+ call_kwargs = mock_init.call_args.kwargs
+ assert call_kwargs.get("storage_backend") == "local"
+
+
+# ---------------------------------------------------------------------------
+# Qdrant optional-dependency guard (now on QdrantVectorStore)
+# ---------------------------------------------------------------------------
+
+class TestQdrantOptionalDepGuard:
+ def test_missing_qdrant_raises_on_open(self, tmp_path):
+ """When qdrant-client is absent, QdrantVectorStore.open() raises RuntimeError."""
+ with patch("foretrieval.vector_store.qdrant._QDRANT_AVAILABLE", False):
+ store = QdrantVectorStore()
+ with pytest.raises(RuntimeError, match="qdrant-client"):
+ store.open("idx", tmp_path, create=True, dim=8)
+
+ def test_missing_qdrant_message_contains_install_hint(self, tmp_path):
+ with patch("foretrieval.vector_store.qdrant._QDRANT_AVAILABLE", False):
+ store = QdrantVectorStore()
+ with pytest.raises(RuntimeError) as exc_info:
+ store.open("idx", tmp_path, create=True, dim=8)
+ assert "foretrieval[qdrant]" in str(exc_info.value)
+
+
+# ---------------------------------------------------------------------------
+# Integration test — full Qdrant index + search (requires GPU + qdrant)
+# ---------------------------------------------------------------------------
+
+@pytest.mark.slow
+@pytest.mark.integration
+def test_qdrant_index_and_search(tmp_path):
+ """Full index → from_index → search cycle using the Qdrant backend.
+
+ Requires:
+ - A compatible GPU (CUDA sm_70+)
+ - qdrant-client installed (pip install foretrieval[qdrant])
+ - test PDFs in tests/data/
+ """
+ if not _QDRANT_AVAILABLE:
+ pytest.skip("qdrant-client not installed; install with: pip install foretrieval[qdrant]")
+
+ from foretrieval import MultiModalRetrieverModel
+
+ retriever = MultiModalRetrieverModel.from_pretrained(
+ pretrained_model_name_or_path="vidore/colqwen2.5-v0.2",
+ index_root=str(tmp_path),
+ storage_backend="qdrant",
+ device=None, # auto-detect
+ verbose=0,
+ )
+
+ retriever.index(
+ input_path=str(DATA_DIR),
+ index_name="qdrant_integration_test",
+ store_collection_with_index=False,
+ overwrite=True,
+ )
+
+ retriever2 = MultiModalRetrieverModel.from_index(
+ index_path="qdrant_integration_test",
+ index_root=str(tmp_path),
+ device=None,
+ )
+
+ results = retriever2.search(
+ "maximum output current", k=1, return_base64_results=False
+ )
+
+ assert len(results) >= 1
+ assert results[0].score is not None
+ assert results[0].doc_id is not None
+ assert results[0].page_num is not None
diff --git a/tests/test_recursive_indexing.py b/tests/test_recursive_indexing.py
new file mode 100644
index 0000000..0f92ba2
--- /dev/null
+++ b/tests/test_recursive_indexing.py
@@ -0,0 +1,280 @@
+"""Tests for recursive directory indexing in ColPaliModel.
+
+Verifies that index(), _process_directory(), and update_index_from_folder()
+traverse subdirectories recursively using rglob instead of iterdir.
+
+No GPU or real model required — ColPaliModel is built with mocked loading
+and _process_and_add_to_index is patched to avoid actual embedding work.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch, call
+
+import pytest
+
+from foretrieval.vector_store import make_vector_store
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_model(tmp_path: Path, storage_backend: str = "local"):
+ """Build a minimal ColPaliModel with GPU/model loading mocked."""
+ from foretrieval.colpali import ColPaliModel
+
+ model = ColPaliModel.__new__(ColPaliModel)
+ model.pretrained_model_name_or_path = "vidore/colpali-v1.2-test"
+ model.model_name = "vidore/colpali-v1.2-test"
+ model.verbose = 0
+ model.load_from_index = False
+ model.index_root = str(tmp_path / "index_root")
+ model.index_name = None
+ model.kwargs = {}
+ model.storage_backend = storage_backend
+ model.storage_config = {}
+ model._storage_qdrant_compat = False
+ model.ingestion = {"backend": "default"}
+ model.ingestion_backend = "default"
+ model.n_gpu = 0
+ model.device = "cpu"
+ model.load_in_4bit = False
+ model.load_in_8bit = False
+ model.bnb_4bit_quant_type = "nf4"
+ model.bnb_4bit_compute_dtype = "float16"
+ model.collection = {}
+ model.embed_id_to_extra = {}
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.doc_ids = set()
+ model.enable_heatmaps = False
+ model.enable_circle = False
+ model.full_document_collection = False
+ model.resize_stored_images = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.docling_dir = None
+ model.SOURCE_EXTS = {".docx", ".txt", ".png", ".jpg"}
+ model.IMAGE_EXTS = {".png", ".jpg", ".jpeg"}
+ model._remote_client = None
+ model.model = None
+ model.processor = MagicMock()
+ model.vector_store = make_vector_store(storage_backend, {})
+ return model
+
+
+def _make_nested_corpus(base: Path) -> list[Path]:
+ """Create a nested directory of dummy files and return their paths sorted."""
+ # base/
+ # a.pdf
+ # sub1/
+ # b.pdf
+ # sub2/
+ # c.pdf
+ (base / "sub1" / "sub2").mkdir(parents=True)
+ files = [
+ base / "a.pdf",
+ base / "sub1" / "b.pdf",
+ base / "sub1" / "sub2" / "c.pdf",
+ ]
+ for f in files:
+ f.write_bytes(b"%PDF-1.4 dummy")
+ return sorted(files, key=lambda p: p.relative_to(base))
+
+
+# ---------------------------------------------------------------------------
+# index() — recursive traversal
+# ---------------------------------------------------------------------------
+
+class TestIndexRecursive:
+ def test_index_visits_all_nested_files(self, tmp_path):
+ """index() must find files in subdirectories via rglob."""
+ corpus = tmp_path / "corpus"
+ corpus.mkdir()
+ expected_files = _make_nested_corpus(corpus)
+
+ model = _make_model(tmp_path)
+ visited = []
+
+ def fake_process(item, *args, **kwargs):
+ visited.append(Path(item))
+ model.doc_ids_to_file_names[kwargs.get("doc_id", len(visited))] = str(item)
+ model.doc_ids.add(kwargs.get("doc_id", len(visited)))
+ model.highest_doc_id = max(model.highest_doc_id, kwargs.get("doc_id", len(visited)))
+ return item
+
+ with (
+ patch.object(model, "_process_and_add_to_index", side_effect=fake_process),
+ patch.object(model, "_export_index"),
+ patch.object(model.vector_store, "open"),
+ patch.object(model.vector_store, "set_processor", create=True),
+ patch.object(model.vector_store, "set_doc_id_to_metadata", create=True),
+ ):
+ model.index(
+ input_path=corpus,
+ index_name="nested_test",
+ overwrite=False,
+ )
+
+ assert sorted(visited, key=lambda p: p.relative_to(corpus)) == expected_files
+
+ def test_index_flat_dir_unchanged_behavior(self, tmp_path):
+ """Flat directory: rglob and old iterdir produce identical file sets."""
+ corpus = tmp_path / "flat"
+ corpus.mkdir()
+ files = []
+ for name in ["x.pdf", "y.pdf", "z.pdf"]:
+ f = corpus / name
+ f.write_bytes(b"%PDF-1.4 dummy")
+ files.append(f)
+ files = sorted(files, key=lambda p: p.relative_to(corpus))
+
+ model = _make_model(tmp_path)
+ visited = []
+
+ def fake_process(item, *args, **kwargs):
+ visited.append(Path(item))
+ model.doc_ids_to_file_names[kwargs.get("doc_id", len(visited))] = str(item)
+ model.doc_ids.add(kwargs.get("doc_id", len(visited)))
+ model.highest_doc_id = max(model.highest_doc_id, kwargs.get("doc_id", len(visited)))
+ return item
+
+ with (
+ patch.object(model, "_process_and_add_to_index", side_effect=fake_process),
+ patch.object(model, "_export_index"),
+ patch.object(model.vector_store, "open"),
+ patch.object(model.vector_store, "set_processor", create=True),
+ patch.object(model.vector_store, "set_doc_id_to_metadata", create=True),
+ ):
+ model.index(input_path=corpus, index_name="flat_test", overwrite=False)
+
+ assert sorted(visited, key=lambda p: p.relative_to(corpus)) == files
+
+ def test_index_overwrite_false_no_existing_index_creates_new(self, tmp_path):
+ """overwrite=False with no existing index directory must proceed normally."""
+ corpus = tmp_path / "corpus"
+ corpus.mkdir()
+ (corpus / "doc.pdf").write_bytes(b"%PDF-1.4 dummy")
+
+ model = _make_model(tmp_path)
+ visited = []
+
+ def fake_process(item, *args, **kwargs):
+ visited.append(Path(item))
+ model.doc_ids_to_file_names[kwargs.get("doc_id", 0)] = str(item)
+ model.doc_ids.add(kwargs.get("doc_id", 0))
+ model.highest_doc_id = max(model.highest_doc_id, kwargs.get("doc_id", 0))
+ return item
+
+ with (
+ patch.object(model, "_process_and_add_to_index", side_effect=fake_process),
+ patch.object(model, "_export_index"),
+ patch.object(model.vector_store, "open"),
+ patch.object(model.vector_store, "set_processor", create=True),
+ patch.object(model.vector_store, "set_doc_id_to_metadata", create=True),
+ ):
+ # Must NOT raise and must NOT return None
+ result = model.index(
+ input_path=corpus,
+ index_name="brand_new_index",
+ overwrite=False,
+ )
+
+ assert result is not None, "index() returned None — did not create a new index"
+ assert len(visited) == 1
+
+
+# ---------------------------------------------------------------------------
+# _process_directory() — recursive traversal
+# ---------------------------------------------------------------------------
+
+class TestProcessDirectoryRecursive:
+ def test_process_directory_visits_nested_files(self, tmp_path):
+ """_process_directory() must recurse into subdirectories."""
+ directory = tmp_path / "docs"
+ directory.mkdir()
+ expected_files = _make_nested_corpus(directory)
+
+ model = _make_model(tmp_path)
+ model.index_name = "test_index"
+ visited = []
+
+ def fake_process(item, *args, **kwargs):
+ visited.append(Path(item))
+ return item
+
+ with patch.object(model, "_process_and_add_to_index", side_effect=fake_process):
+ model._process_directory(
+ directory=directory,
+ store_collection_with_index=False,
+ base_doc_id=0,
+ metadata=None,
+ batch_size=1,
+ )
+
+ assert sorted(visited, key=lambda p: p.relative_to(directory)) == expected_files
+
+
+# ---------------------------------------------------------------------------
+# update_index_from_folder() — recursive traversal
+# ---------------------------------------------------------------------------
+
+class TestUpdateIndexFromFolderRecursive:
+ def test_update_visits_nested_new_files(self, tmp_path):
+ """update_index_from_folder() must recurse into subdirectories."""
+ folder = tmp_path / "docs"
+ folder.mkdir()
+ expected_files = _make_nested_corpus(folder)
+
+ model = _make_model(tmp_path)
+ model.index_name = "test_index"
+ model.doc_ids_to_file_names = {}
+ model.doc_ids = set()
+ model.highest_doc_id = -1
+ visited = []
+
+ def fake_process(item, *args, **kwargs):
+ visited.append(Path(item))
+ model.doc_ids_to_file_names[kwargs.get("doc_id", len(visited))] = str(item)
+ model.doc_ids.add(kwargs.get("doc_id", len(visited)))
+ model.highest_doc_id = max(model.highest_doc_id, kwargs.get("doc_id", len(visited)))
+ return item
+
+ with (
+ patch.object(model, "_process_and_add_to_index", side_effect=fake_process),
+ patch.object(model, "_export_index"),
+ ):
+ model.update_index_from_folder(folder=folder)
+
+ assert sorted(visited, key=lambda p: p.relative_to(folder)) == expected_files
+
+ def test_update_skips_already_indexed_nested_files(self, tmp_path):
+ """update_index_from_folder() must skip already-indexed files in subdirs."""
+ folder = tmp_path / "docs"
+ folder.mkdir()
+ (folder / "sub").mkdir()
+ already_indexed = folder / "sub" / "old.pdf"
+ new_file = folder / "sub" / "new.pdf"
+ already_indexed.write_bytes(b"%PDF-1.4 old")
+ new_file.write_bytes(b"%PDF-1.4 new")
+
+ model = _make_model(tmp_path)
+ model.index_name = "test_index"
+ model.doc_ids_to_file_names = {0: str(already_indexed.resolve())}
+ model.doc_ids = {0}
+ model.highest_doc_id = 0
+ visited = []
+
+ def fake_process(item, *args, **kwargs):
+ visited.append(Path(item))
+ return item
+
+ with (
+ patch.object(model, "_process_and_add_to_index", side_effect=fake_process),
+ patch.object(model, "_export_index"),
+ ):
+ model.update_index_from_folder(folder=folder)
+
+ assert visited == [new_file], f"Expected only new.pdf, got {visited}"
diff --git a/tests/test_ssh_utils.py b/tests/test_ssh_utils.py
new file mode 100644
index 0000000..f559273
--- /dev/null
+++ b/tests/test_ssh_utils.py
@@ -0,0 +1,182 @@
+"""Tests for foretrieval.ssh_utils.open_ssh_client.
+
+All tests mock paramiko.SSHClient.connect so nothing touches the network.
+The goal is to confirm that ~/.ssh/config directives are honoured the
+same way as the OpenSSH CLI.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+# ---------------------------------------------------------------------------
+# Fixtures / helpers
+# ---------------------------------------------------------------------------
+
+def _write_ssh_config(tmp_path: Path, content: str) -> Path:
+ """Write an SSH config to a temp file and return its path."""
+ p = tmp_path / "ssh_config"
+ p.write_text(content)
+ return p
+
+
+def _captured_connect_kwargs(mock_client_class):
+ """Return the kwargs of the most recent SSHClient.connect call."""
+ instance = mock_client_class.return_value
+ return instance.connect.call_args.kwargs
+
+
+# ---------------------------------------------------------------------------
+# Basic resolution
+# ---------------------------------------------------------------------------
+
+class TestHostAliasResolution:
+ def test_alias_resolves_to_hostname(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, """
+Host pf01
+ HostName pf01.example.com
+ User alice
+ Port 2222
+""")
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient:
+ ssh_utils.open_ssh_client("pf01", ssh_config_path=cfg)
+ kw = _captured_connect_kwargs(MockClient)
+ assert kw["hostname"] == "pf01.example.com"
+ assert kw["port"] == 2222
+ assert kw["username"] == "alice"
+
+ def test_no_config_falls_back_to_direct(self, tmp_path):
+ """Missing config file → direct connect with the raw host string."""
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient:
+ ssh_utils.open_ssh_client(
+ "real.example.com", ssh_config_path=tmp_path / "missing"
+ )
+ kw = _captured_connect_kwargs(MockClient)
+ assert kw["hostname"] == "real.example.com"
+ assert kw["port"] == 22
+
+ def test_no_match_falls_back_to_direct(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, """
+Host someotheralias
+ HostName irrelevant
+""")
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient:
+ ssh_utils.open_ssh_client("pf01", ssh_config_path=cfg)
+ kw = _captured_connect_kwargs(MockClient)
+ assert kw["hostname"] == "pf01"
+
+
+# ---------------------------------------------------------------------------
+# Override precedence
+# ---------------------------------------------------------------------------
+
+class TestOverrides:
+ def test_user_override_beats_ssh_config(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, """
+Host pf01
+ HostName pf01.example.com
+ User alice
+""")
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient:
+ ssh_utils.open_ssh_client(
+ "pf01", ssh_user="override_user", ssh_config_path=cfg,
+ )
+ kw = _captured_connect_kwargs(MockClient)
+ assert kw["username"] == "override_user"
+
+ def test_key_path_override_beats_identityfile(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, """
+Host pf01
+ HostName pf01.example.com
+ IdentityFile ~/.ssh/should_be_ignored
+""")
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient:
+ ssh_utils.open_ssh_client(
+ "pf01",
+ ssh_key_path="/explicit/key/path",
+ ssh_config_path=cfg,
+ )
+ kw = _captured_connect_kwargs(MockClient)
+ assert kw["key_filename"] == "/explicit/key/path"
+
+
+# ---------------------------------------------------------------------------
+# ProxyCommand / ProxyJump
+# ---------------------------------------------------------------------------
+
+class TestProxy:
+ def test_proxycommand_creates_sock(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, """
+Host pf01
+ HostName pf01.example.com
+ User alice
+ ProxyCommand ssh -W %h:%p bastion
+""")
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient, \
+ patch("paramiko.ProxyCommand") as MockPC:
+ ssh_utils.open_ssh_client("pf01", ssh_config_path=cfg)
+ # ProxyCommand instance was constructed and passed as 'sock'
+ kw = _captured_connect_kwargs(MockClient)
+ assert "sock" in kw
+ # The expanded command was used
+ MockPC.assert_called_once()
+ expanded_cmd = MockPC.call_args.args[0]
+ assert "pf01.example.com" in expanded_cmd
+ assert ":22" in expanded_cmd
+ assert "bastion" in expanded_cmd
+
+ def test_proxyjump_translates_to_proxycommand(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, """
+Host pf01
+ HostName pf01.example.com
+ ProxyJump bastion.example.com
+""")
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient, \
+ patch("paramiko.ProxyCommand") as MockPC:
+ ssh_utils.open_ssh_client("pf01", ssh_config_path=cfg)
+ kw = _captured_connect_kwargs(MockClient)
+ assert "sock" in kw
+ MockPC.assert_called_once()
+ cmd = MockPC.call_args.args[0]
+ assert "ssh -W" in cmd
+ assert "bastion.example.com" in cmd
+
+ def test_no_proxy_no_sock(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, """
+Host pf01
+ HostName pf01.example.com
+""")
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHClient") as MockClient:
+ ssh_utils.open_ssh_client("pf01", ssh_config_path=cfg)
+ kw = _captured_connect_kwargs(MockClient)
+ assert "sock" not in kw
+
+
+# ---------------------------------------------------------------------------
+# Failure modes
+# ---------------------------------------------------------------------------
+
+class TestFailureModes:
+ def test_malformed_ssh_config_falls_back_to_direct(self, tmp_path):
+ cfg = _write_ssh_config(tmp_path, "this is not valid ssh config syntax{{{")
+ # paramiko.SSHConfig is permissive — write something that *does*
+ # crash its parser. Easiest: simulate the parse exception.
+ from foretrieval import ssh_utils
+ with patch("paramiko.SSHConfig") as MockCfg, \
+ patch("paramiko.SSHClient") as MockClient:
+ MockCfg.return_value.parse.side_effect = RuntimeError("bad syntax")
+ ssh_utils.open_ssh_client("pf01", ssh_config_path=cfg)
+ kw = _captured_connect_kwargs(MockClient)
+ assert kw["hostname"] == "pf01"
diff --git a/tests/test_task10_bugfixes.py b/tests/test_task10_bugfixes.py
new file mode 100644
index 0000000..df7df6b
--- /dev/null
+++ b/tests/test_task10_bugfixes.py
@@ -0,0 +1,296 @@
+"""Regression tests for task_10 bug fixes.
+
+Bug 1: build_metadata_list_for_dir recursive alignment — covered in test_metadata_no_ai.py
+Bug 2: cleanup on index failure
+Bug 3: doc_ids/highest_doc_id derived from doc_ids_to_file_names (not doc_id_to_metadata)
+Bug 4: matplotlib.colormaps replaces deprecated cm.get_cmap
+"""
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_model(storage_backend: str = "local", index_root: str = "/tmp/foretrieval_test"):
+ """Minimal ColPaliModel with GPU / model loading patched out."""
+ with (
+ patch("foretrieval.colpali.ColPaliModel._load_model_and_processor"),
+ patch("foretrieval.colpali.ColPaliModel._load_processor_only"),
+ ):
+ from foretrieval.colpali import ColPaliModel
+ from foretrieval.vector_store.local import LocalVectorStore
+
+ model = ColPaliModel.__new__(ColPaliModel)
+ model.pretrained_model_name_or_path = "vidore/colpali-v1.2-test"
+ model.model_name = "vidore/colpali-v1.2-test"
+ model.verbose = 0
+ model.load_from_index = False
+ model.index_root = index_root
+ model.index_name = None
+ model.kwargs = {}
+ model.storage_backend = storage_backend
+ model.storage_config = {}
+ model._storage_qdrant_compat = False
+ model.ingestion = {"backend": "default"}
+ model.ingestion_backend = "default"
+ model.n_gpu = 0
+ model.device = "cpu"
+ model.load_in_4bit = False
+ model.load_in_8bit = False
+ model.bnb_4bit_quant_type = "nf4"
+ model.bnb_4bit_compute_dtype = "float16"
+ model.collection = {}
+ model.embed_id_to_extra = {}
+ model.doc_id_to_metadata = {}
+ model.doc_ids_to_file_names = {}
+ model.doc_ids = set()
+ model.enable_heatmaps = False
+ model.enable_circle = False
+ model.full_document_collection = False
+ model.resize_stored_images = False
+ model.max_image_width = None
+ model.max_image_height = None
+ model.highest_doc_id = -1
+ model.docling_dir = None
+ model.SOURCE_EXTS = set()
+ model.IMAGE_EXTS = set()
+ model._remote_client = None
+ model.model = None
+ model.processor = MagicMock()
+ model.vector_store = LocalVectorStore()
+ model.index_description = ""
+ return model
+
+
+# ---------------------------------------------------------------------------
+# Bug 3: load without metadata yields correct doc_ids / highest_doc_id
+# ---------------------------------------------------------------------------
+
+class TestLoadWithoutMetadata:
+ """Regression: doc_ids / highest_doc_id are derived from doc_ids_to_file_names."""
+
+ def _build_index_dir(self, tmp_path: Path, with_metadata: bool) -> Path:
+ """Write minimal index sidecars to tmp_path."""
+ import srsly, torch
+
+ idx = tmp_path / "myindex"
+ idx.mkdir()
+
+ # index_config.json.gz
+ srsly.write_gzip_json(idx / "index_config.json.gz", {
+ "model_name": "vidore/colpali-v1.2-test",
+ "full_document_collection": False,
+ "highest_doc_id": 2,
+ "resize_stored_images": False,
+ "max_image_width": None,
+ "max_image_height": None,
+ "library_version": "0.0.0",
+ "storage_backend": "local",
+ "storage_config": None,
+ "description": "",
+ })
+
+ # doc_ids_to_file_names.json.gz — ALWAYS present
+ srsly.write_gzip_json(idx / "doc_ids_to_file_names.json.gz", {
+ "0": "/data/doc_a.pdf",
+ "1": "/data/doc_b.pdf",
+ "2": "/data/doc_c.pdf",
+ })
+
+ # embed_id_to_doc_id.json.gz
+ srsly.write_gzip_json(idx / "embed_id_to_doc_id.json.gz", {})
+
+ # embed_id_to_extra.pt
+ torch.save({}, idx / "embed_id_to_extra.pt")
+
+ if with_metadata:
+ srsly.write_gzip_json(idx / "metadata.json.gz", {
+ "0": {"stem": "doc_a", "ext": ".pdf"},
+ "1": {"stem": "doc_b", "ext": ".pdf"},
+ "2": {"stem": "doc_c", "ext": ".pdf"},
+ })
+ # No metadata.json.gz when with_metadata=False
+
+ return idx
+
+ def test_load_without_metadata_doc_ids_correct(self, tmp_path):
+ """doc_ids is populated from doc_ids_to_file_names when metadata absent."""
+ idx = self._build_index_dir(tmp_path, with_metadata=False)
+ model = _make_model(index_root=str(tmp_path))
+
+ model._load_local_sidecars(idx)
+
+ # Simulate _load_index_state final derivation
+ id_source = model.doc_ids_to_file_names or model.doc_id_to_metadata
+ model.highest_doc_id = max(id_source.keys(), default=-1)
+ model.doc_ids = set(id_source.keys())
+
+ assert model.doc_ids == {0, 1, 2}, (
+ f"Expected {{0, 1, 2}}, got {model.doc_ids}"
+ )
+ assert model.highest_doc_id == 2, (
+ f"Expected 2, got {model.highest_doc_id}"
+ )
+ assert model.doc_id_to_metadata == {}, "Metadata should be empty"
+
+ def test_load_with_metadata_doc_ids_correct(self, tmp_path):
+ """doc_ids is populated correctly when metadata present."""
+ idx = self._build_index_dir(tmp_path, with_metadata=True)
+ model = _make_model(index_root=str(tmp_path))
+
+ model._load_local_sidecars(idx)
+
+ id_source = model.doc_ids_to_file_names or model.doc_id_to_metadata
+ model.highest_doc_id = max(id_source.keys(), default=-1)
+ model.doc_ids = set(id_source.keys())
+
+ assert model.doc_ids == {0, 1, 2}
+ assert model.highest_doc_id == 2
+
+ def test_load_local_sidecars_no_metadata(self, tmp_path):
+ """_load_local_sidecars: doc_id_to_metadata is {} and doc_ids_to_file_names populated."""
+ idx = self._build_index_dir(tmp_path, with_metadata=False)
+ model = _make_model(index_root=str(tmp_path))
+
+ model._load_local_sidecars(idx)
+
+ assert model.doc_id_to_metadata == {}
+ assert set(model.doc_ids_to_file_names.keys()) == {0, 1, 2}
+
+ def test_apply_bookkeeping_blob_without_metadata(self):
+ """Remote path: _apply_bookkeeping_blob derives doc_ids from doc_ids_to_file_names."""
+ model = _make_model(storage_backend="remote")
+
+ blob = {
+ "index_config": {
+ "model_name": "vidore/colpali-v1.2-test",
+ "full_document_collection": False,
+ "highest_doc_id": 2,
+ "resize_stored_images": False,
+ "max_image_width": None,
+ "max_image_height": None,
+ "description": "",
+ },
+ "embed_id_to_extra": {},
+ "doc_ids_to_file_names": {
+ "0": "/data/doc_a.pdf",
+ "1": "/data/doc_b.pdf",
+ "2": "/data/doc_c.pdf",
+ },
+ "doc_id_to_metadata": {}, # empty: index built without add_metadata
+ }
+
+ model._apply_bookkeeping_blob(blob)
+
+ assert model.doc_ids == {0, 1, 2}, (
+ f"Expected {{0, 1, 2}}, got {model.doc_ids}"
+ )
+ assert model.highest_doc_id == 2
+
+
+# ---------------------------------------------------------------------------
+# Bug 2: cleanup on index failure
+# ---------------------------------------------------------------------------
+
+class TestIndexCleanup:
+ """Regression: _cleanup_failed_index removes partial artefacts."""
+
+ def test_cleanup_removes_local_index_dir(self, tmp_path):
+ """_cleanup_failed_index removes the local index directory."""
+ index_name = "partial_index"
+ index_root = tmp_path
+ index_path = index_root / index_name
+ index_path.mkdir()
+ (index_path / "index_config.json.gz").write_bytes(b"fake")
+
+ model = _make_model(index_root=str(index_root))
+ model.index_name = index_name
+ model.highest_doc_id = 0
+ model.doc_ids = {0}
+
+ model._cleanup_failed_index(index_name)
+
+ assert not index_path.exists(), "Partial index directory should be removed"
+ assert model.index_name is None
+ assert model.highest_doc_id == -1
+ assert model.doc_ids == set()
+
+ def test_cleanup_tolerates_missing_dir(self, tmp_path):
+ """_cleanup_failed_index does not raise if the directory doesn't exist."""
+ model = _make_model(index_root=str(tmp_path))
+ # Should not raise
+ model._cleanup_failed_index("nonexistent_index")
+
+ def test_index_raises_cleans_up_on_metadata_mismatch(self, tmp_path):
+ """index() cleans up when metadata list length doesn't match file count."""
+ from foretrieval.models_metadata import DocMetadata
+ from foretrieval.vector_store.local import LocalVectorStore
+
+ # Create a nested data directory with 2 files
+ data_dir = tmp_path / "docs"
+ data_dir.mkdir()
+ (data_dir / "a.pdf").write_bytes(b"%PDF")
+ subdir = data_dir / "sub"
+ subdir.mkdir()
+ (subdir / "b.pdf").write_bytes(b"%PDF")
+
+ model = _make_model(index_root=str(tmp_path))
+ model.vector_store = LocalVectorStore()
+
+ # Provide metadata with wrong length (1 instead of 2)
+ wrong_metadata = [DocMetadata(stem="only_one", ext=".pdf")]
+
+ with pytest.raises(ValueError, match="metadata entries"):
+ model.index(
+ input_path=str(data_dir),
+ index_name="test_cleanup",
+ metadata=wrong_metadata,
+ )
+
+ # Index directory should be cleaned up
+ index_path = tmp_path / "test_cleanup"
+ assert not index_path.exists(), (
+ "Partial index directory should be removed after failed indexing"
+ )
+
+
+# ---------------------------------------------------------------------------
+# Bug 4: matplotlib.colormaps replaces deprecated cm.get_cmap
+# ---------------------------------------------------------------------------
+
+class TestHeatmapColormap:
+ """Regression: heatmap_overlay_base64 works with matplotlib >= 3.9."""
+
+ def test_heatmap_overlay_base64_returns_string(self):
+ """heatmap_overlay_base64 returns a base64 string without AttributeError."""
+ import numpy as np
+ from PIL import Image
+ from foretrieval.plot_utils import heatmap_overlay_base64
+
+ img = Image.fromarray(
+ (np.ones((32, 32, 3)) * 200).astype("uint8")
+ )
+ heat = (np.ones((32, 32)) * 128).astype("uint8")
+ result = heatmap_overlay_base64(img, heat, cmap="jet")
+ assert isinstance(result, str) and len(result) > 0
+
+ def test_heatmap_overlay_base64_different_cmaps(self):
+ """Several colormaps all work without AttributeError."""
+ import numpy as np
+ from PIL import Image
+ from foretrieval.plot_utils import heatmap_overlay_base64
+
+ img = Image.fromarray(
+ (np.ones((16, 16, 3)) * 100).astype("uint8")
+ )
+ heat = (np.ones((16, 16)) * 50).astype("uint8")
+ for cmap in ("jet", "viridis", "plasma", "hot"):
+ result = heatmap_overlay_base64(img, heat, cmap=cmap)
+ assert isinstance(result, str), f"cmap={cmap} returned non-string"
diff --git a/tests/test_task11_features.py b/tests/test_task11_features.py
new file mode 100644
index 0000000..1a24eda
--- /dev/null
+++ b/tests/test_task11_features.py
@@ -0,0 +1,130 @@
+"""Unit tests for task_11 Feature 3: circle center alignment with heatmap.
+
+Verifies that draw_circle_on_max_patch applies patch growth before argmax
+so the circle center aligns with the same heat region as heatmap_overlay_base64.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import pytest
+from PIL import Image
+
+from foretrieval.plot_utils import (
+ draw_circle_on_max_patch,
+ grow_heatmap_patches_torch,
+ heatmap_overlay_base64,
+)
+
+
+def _synthetic_image(w: int = 100, h: int = 80) -> Image.Image:
+ """Create a plain grey synthetic image."""
+ arr = np.full((h, w, 3), 128, dtype=np.uint8)
+ return Image.fromarray(arr, "RGB")
+
+
+def _heat_with_peak(Hp: int, Wp: int, peak_r: int, peak_c: int) -> torch.Tensor:
+ """Create a heat grid with a clear peak at (peak_r, peak_c)."""
+ heat = torch.zeros(Hp, Wp)
+ heat[peak_r, peak_c] = 10.0
+ return heat
+
+
+class TestCircleGrowthAlignment:
+ """draw_circle_on_max_patch with patch_grow_pct=300,grow_mode='mean'
+ should find the same argmax as the grown heat used by heatmap_overlay_base64.
+ """
+
+ def test_grow_heatmap_patches_torch_no_op_at_100(self):
+ heat = torch.tensor([[1.0, 0.5], [0.2, 0.8]])
+ out = grow_heatmap_patches_torch(heat, patch_grow_pct=100.0)
+ assert torch.allclose(out, heat)
+
+ def test_grow_heatmap_patches_torch_grow_mean(self):
+ heat = torch.zeros(5, 5)
+ heat[2, 2] = 10.0
+ grown = grow_heatmap_patches_torch(heat, patch_grow_pct=300.0, grow_mode="mean")
+ # 300% → radius=2 → 5×5 avg_pool. Peak should spread.
+ assert grown.shape == (5, 5)
+ assert grown.max() < 10.0 # mean dilutes the peak
+ assert grown[2, 2] > 0.0 # center still non-zero
+
+ def test_circle_argmax_matches_grown_heat_argmax(self):
+ """When peak is clear, circle should center on same patch as heatmap peak."""
+ Hp, Wp = 8, 8
+ # Put strong peak at (3,5)
+ heat = _heat_with_peak(Hp, Wp, peak_r=3, peak_c=5)
+
+ # Grown heat argmax
+ grown = grow_heatmap_patches_torch(heat.clone(), patch_grow_pct=300.0, grow_mode="mean")
+ flat_grown = int(torch.argmax(grown.flatten()))
+ r_grown = flat_grown // Wp
+ c_grown = flat_grown % Wp
+
+ # draw_circle_on_max_patch with same grow
+ img = _synthetic_image()
+ W, H = img.size
+ patch_w = W / float(Wp)
+ patch_h = H / float(Hp)
+ expected_cx = (c_grown + 0.5) * patch_w
+ expected_cy = (r_grown + 0.5) * patch_h
+
+ # Verify grow function defined before draw_circle_on_max_patch (import order)
+ from foretrieval import plot_utils
+ import inspect
+ src = inspect.getsource(plot_utils)
+ grow_pos = src.index("def grow_heatmap_patches_torch")
+ circle_pos = src.index("def draw_circle_on_max_patch")
+ assert grow_pos < circle_pos, "grow_heatmap_patches_torch must be defined before draw_circle_on_max_patch"
+
+ def test_draw_circle_returns_rgb_image(self):
+ img = _synthetic_image()
+ heat = _heat_with_peak(8, 8, 3, 5)
+ result = draw_circle_on_max_patch(img, heat, patch_grow_pct=300.0, grow_mode="mean")
+ assert isinstance(result, Image.Image)
+ assert result.mode == "RGB"
+ assert result.size == img.size
+
+ def test_draw_circle_no_grow_still_works(self):
+ """Default grow (100%) should work as before."""
+ img = _synthetic_image()
+ heat = _heat_with_peak(8, 8, 3, 5)
+ result = draw_circle_on_max_patch(img, heat)
+ assert isinstance(result, Image.Image)
+ assert result.mode == "RGB"
+
+ def test_draw_circle_modified_image(self):
+ """Circle is actually drawn — result differs from input."""
+ img = _synthetic_image()
+ heat = _heat_with_peak(8, 8, 3, 5)
+ result = draw_circle_on_max_patch(img, heat, patch_grow_pct=300.0, grow_mode="mean")
+ assert np.array(result).shape == np.array(img).shape
+ # Image should differ (circle pixels added)
+ diff = np.abs(np.array(result).astype(float) - np.array(img).astype(float))
+ assert diff.max() > 0, "Result should differ from input (circle drawn)"
+
+
+class TestHeatmapAndCircleConsistency:
+ """Ensure circle and heatmap use the same heat processing path."""
+
+ def test_heatmap_overlay_returns_nonempty_base64(self):
+ img = _synthetic_image()
+ heat = _heat_with_peak(8, 8, 3, 5)
+ b64 = heatmap_overlay_base64(
+ img, heat, patch_grow_pct=300.0, grow_mode="mean"
+ )
+ assert isinstance(b64, str)
+ assert len(b64) > 100
+
+ def test_colpali_circle_call_uses_grow_params(self):
+ """Smoke-test that the colpali call site now passes grow params.
+
+ We can't import ColPaliModel without GPU weights, but we can verify
+ the source code of colpali.py contains the new call signature.
+ """
+ from pathlib import Path
+ colpali_src = (Path(__file__).parent.parent / "foretrieval" / "colpali.py").read_text()
+ assert "patch_grow_pct=300.0" in colpali_src, \
+ "colpali.py circle call site should pass patch_grow_pct=300.0"
+ assert "grow_mode=\"mean\"" in colpali_src, \
+ "colpali.py circle call site should pass grow_mode='mean'"
diff --git a/tests/test_vector_db_server_admin.py b/tests/test_vector_db_server_admin.py
new file mode 100644
index 0000000..5b4a698
--- /dev/null
+++ b/tests/test_vector_db_server_admin.py
@@ -0,0 +1,266 @@
+"""Tests for the vector-DB server admin endpoints.
+
+Covers:
+- /v1/admin/indexes
+- /v1/admin/data_folders
+- 60 s TTL cache + invalidation on data-plane writes
+- Auth gating (inherits the shared middleware)
+- Symlink jail / containment
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import time
+from pathlib import Path
+
+import pytest
+from fastapi.testclient import TestClient
+
+# Patch FOR_DB_DATA_DIR before importing the server (mirrors the existing
+# test_vector_db_server_app.py pattern).
+_TMP = tempfile.mkdtemp(prefix="foretrieval_db_admin_test_")
+
+import foretrieval.vector_db_server.server as _server_mod
+
+_server_mod._DATA_DIR = Path(_TMP)
+_server_mod._API_KEY = None
+_server_mod._registry.clear()
+_server_mod._locks.clear()
+_server_mod._SIZE_CACHE.clear()
+
+from foretrieval.vector_db_server.server import app # noqa: E402
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _mk_index_dir(name: str, n_files: int = 3, size_per_file: int = 100) -> Path:
+ """Create a fake index directory under _DATA_DIR."""
+ p = _server_mod._DATA_DIR / name
+ p.mkdir(parents=True, exist_ok=True)
+ # Mark as an index by adding the server-side sentinel (written by _write_meta).
+ (p / "index.json").write_text('{"backend": "qdrant", "storage_config": null}')
+ for i in range(n_files):
+ (p / f"data_{i}.bin").write_bytes(b"x" * size_per_file)
+ return p
+
+
+def _mk_data_dir(name: str, n_files: int = 2, size_per_file: int = 50) -> Path:
+ """Create a plain folder (NOT an index) under _DATA_DIR."""
+ p = _server_mod._DATA_DIR / name
+ p.mkdir(parents=True, exist_ok=True)
+ for i in range(n_files):
+ (p / f"file_{i}.bin").write_bytes(b"y" * size_per_file)
+ return p
+
+
+@pytest.fixture(autouse=True)
+def _clean_state():
+ _server_mod._registry.clear()
+ _server_mod._locks.clear()
+ _server_mod._SIZE_CACHE.clear()
+ # Wipe data dir between tests
+ if _server_mod._DATA_DIR.exists():
+ for entry in list(_server_mod._DATA_DIR.iterdir()):
+ if entry.is_dir():
+ import shutil
+ shutil.rmtree(entry, ignore_errors=True)
+ else:
+ entry.unlink(missing_ok=True)
+ yield
+ _server_mod._registry.clear()
+ _server_mod._locks.clear()
+ _server_mod._SIZE_CACHE.clear()
+
+
+@pytest.fixture()
+def client():
+ return TestClient(app)
+
+
+# ---------------------------------------------------------------------------
+# /v1/admin/indexes
+# ---------------------------------------------------------------------------
+
+class TestAdminIndexes:
+ def test_empty_data_dir(self, client):
+ # Nothing under data_dir
+ r = client.get("/v1/admin/indexes")
+ assert r.status_code == 200
+ payload = r.json()
+ assert payload["items"] == []
+ assert payload["count"] == 0
+ assert payload["data_dir"] == str(_server_mod._DATA_DIR)
+
+ def test_lists_index_dirs(self, client):
+ _mk_index_dir("alpha", n_files=4, size_per_file=200)
+ _mk_index_dir("beta", n_files=1, size_per_file=10)
+ # A data-only folder should NOT appear in /indexes
+ _mk_data_dir("gamma_data")
+
+ r = client.get("/v1/admin/indexes")
+ assert r.status_code == 200
+ items = {it["name"]: it for it in r.json()["items"]}
+ assert set(items) == {"alpha", "beta"}
+ assert items["alpha"]["n_files"] >= 4 # data files + index_config marker
+ assert items["alpha"]["size_bytes"] > 0
+ assert items["beta"]["size_bytes"] > 0
+ assert "modified" in items["alpha"]
+
+
+# ---------------------------------------------------------------------------
+# /v1/admin/data_folders
+# ---------------------------------------------------------------------------
+
+class TestAdminDataFolders:
+ def test_lists_every_subdir_with_is_index_flag(self, client):
+ _mk_index_dir("alpha")
+ _mk_data_dir("uploads")
+
+ r = client.get("/v1/admin/data_folders")
+ assert r.status_code == 200
+ items = {it["name"]: it for it in r.json()["items"]}
+ assert set(items) == {"alpha", "uploads"}
+ assert items["alpha"]["is_index"] is True
+ assert items["uploads"]["is_index"] is False
+
+ def test_empty(self, client):
+ r = client.get("/v1/admin/data_folders")
+ assert r.status_code == 200
+ assert r.json()["count"] == 0
+
+
+# ---------------------------------------------------------------------------
+# Auth gating
+# ---------------------------------------------------------------------------
+
+class TestAdminAuth:
+ def test_no_token_no_auth(self, client):
+ _mk_index_dir("alpha")
+ # _API_KEY is None — no auth required
+ r = client.get("/v1/admin/indexes")
+ assert r.status_code == 200
+
+ def test_with_token_missing_header(self, client):
+ _server_mod._API_KEY = "secret"
+ try:
+ r = client.get("/v1/admin/indexes")
+ assert r.status_code == 401
+ finally:
+ _server_mod._API_KEY = None
+
+ def test_with_token_wrong(self, client):
+ _server_mod._API_KEY = "secret"
+ try:
+ r = client.get(
+ "/v1/admin/indexes",
+ headers={"Authorization": "Bearer wrong"},
+ )
+ assert r.status_code == 401
+ finally:
+ _server_mod._API_KEY = None
+
+ def test_with_token_correct(self, client):
+ _mk_index_dir("alpha")
+ _server_mod._API_KEY = "secret"
+ try:
+ r = client.get(
+ "/v1/admin/indexes",
+ headers={"Authorization": "Bearer secret"},
+ )
+ assert r.status_code == 200
+ assert any(it["name"] == "alpha" for it in r.json()["items"])
+ finally:
+ _server_mod._API_KEY = None
+
+
+# ---------------------------------------------------------------------------
+# TTL cache + invalidation
+# ---------------------------------------------------------------------------
+
+class TestSizeCache:
+ def test_cache_hit_within_ttl(self, client, monkeypatch):
+ _mk_index_dir("alpha", n_files=2, size_per_file=100)
+
+ call_count = {"n": 0}
+ real = _server_mod._dir_stats_sync
+
+ def spy(path):
+ call_count["n"] += 1
+ return real(path)
+
+ monkeypatch.setattr(_server_mod, "_dir_stats_sync", spy)
+
+ client.get("/v1/admin/indexes")
+ first = call_count["n"]
+ client.get("/v1/admin/indexes")
+ second = call_count["n"]
+
+ # Second call must hit the cache — no additional walks
+ assert second == first
+
+ def test_cache_miss_after_ttl_expiry(self, client, monkeypatch):
+ _mk_index_dir("alpha", n_files=2, size_per_file=100)
+ monkeypatch.setattr(_server_mod, "_CACHE_TTL", 0.0)
+
+ call_count = {"n": 0}
+ real = _server_mod._dir_stats_sync
+
+ def spy(path):
+ call_count["n"] += 1
+ return real(path)
+
+ monkeypatch.setattr(_server_mod, "_dir_stats_sync", spy)
+
+ client.get("/v1/admin/indexes")
+ first = call_count["n"]
+ client.get("/v1/admin/indexes")
+ second = call_count["n"]
+
+ # TTL is 0 → every request must re-walk
+ assert second == 2 * first
+
+ def test_invalidation_on_delete(self, client):
+ # Manually populate cache then trigger DELETE on a (non-existent) collection;
+ # the cache for that path should be cleared.
+ idx = _mk_index_dir("alpha", n_files=2, size_per_file=100)
+ client.get("/v1/admin/indexes")
+ assert str(idx) in _server_mod._SIZE_CACHE
+ # The delete endpoint always invalidates and returns {"deleted": True}.
+ r = client.delete("/v1/collection/alpha")
+ assert r.status_code == 200
+ assert str(idx) not in _server_mod._SIZE_CACHE
+
+
+# ---------------------------------------------------------------------------
+# Symlink jail
+# ---------------------------------------------------------------------------
+
+class TestSymlinkContainment:
+ def test_symlinked_dir_outside_data_dir_is_skipped(self, client, tmp_path):
+ # Create a real dir OUTSIDE _DATA_DIR
+ outside = tmp_path / "outside"
+ outside.mkdir()
+ (outside / "evil.bin").write_bytes(b"z" * 10000)
+
+ # Symlink INSIDE _DATA_DIR pointing outside
+ link = _server_mod._DATA_DIR / "linked"
+ try:
+ os.symlink(outside, link)
+ except (NotImplementedError, OSError):
+ pytest.skip("symlinks not supported on this platform")
+
+ # Also a legitimate index alongside
+ _mk_index_dir("real_index")
+
+ r = client.get("/v1/admin/data_folders")
+ assert r.status_code == 200
+ names = {it["name"] for it in r.json()["items"]}
+ # The symlink target's content must not be reported via the link.
+ # The directory entry itself may appear, but with size_bytes == 0 because
+ # we resolve and check containment in the walk.
+ assert "real_index" in names
diff --git a/tests/test_vector_db_server_app.py b/tests/test_vector_db_server_app.py
new file mode 100644
index 0000000..6fa5367
--- /dev/null
+++ b/tests/test_vector_db_server_app.py
@@ -0,0 +1,318 @@
+"""Tests for the FastAPI vector-DB server app.
+
+Uses FastAPI's TestClient for in-process HTTP calls — no real server needed.
+All storage is local (temp dir) to keep tests fast and self-contained.
+"""
+
+import os
+import tempfile
+from unittest.mock import patch
+
+import pytest
+import torch
+from fastapi.testclient import TestClient
+
+# ── patch FOR_DB_DATA_DIR before importing the server ──────────────────────
+_TMP = tempfile.mkdtemp(prefix="foretrieval_db_test_")
+
+# We patch the module-level constants before the app is created.
+import foretrieval.vector_db_server.server as _server_mod
+
+_server_mod._DATA_DIR = __import__("pathlib").Path(_TMP)
+_server_mod._API_KEY = None # auth off by default
+_server_mod._registry.clear()
+_server_mod._locks.clear()
+
+from foretrieval.vector_db_server.server import app # noqa: E402
+from foretrieval.vector_db_server.client import _dumps, _loads
+from foretrieval.vector_store.base import make_point_id
+
+
+@pytest.fixture(autouse=True)
+def _clean_registry():
+ """Clear in-memory registry and lock state between tests."""
+ _server_mod._registry.clear()
+ _server_mod._locks.clear()
+ yield
+ _server_mod._registry.clear()
+ _server_mod._locks.clear()
+
+
+@pytest.fixture()
+def client():
+ return TestClient(app)
+
+
+# ---------------------------------------------------------------------------
+# /health
+# ---------------------------------------------------------------------------
+
+class TestHealth:
+ def test_health_ok(self, client):
+ r = client.get("/health")
+ assert r.status_code == 200
+ assert r.json()["status"] == "ok"
+
+
+# ---------------------------------------------------------------------------
+# Auth middleware
+# ---------------------------------------------------------------------------
+
+class TestAuth:
+ def test_missing_auth_header_returns_401(self, client):
+ _server_mod._API_KEY = "secret"
+ try:
+ r = client.get("/v1/collection/myidx/exists")
+ assert r.status_code == 401
+ finally:
+ _server_mod._API_KEY = None
+
+ def test_wrong_token_returns_401(self, client):
+ _server_mod._API_KEY = "secret"
+ try:
+ r = client.get(
+ "/v1/collection/myidx/exists",
+ headers={"Authorization": "Bearer wrongtoken"},
+ )
+ assert r.status_code == 401
+ finally:
+ _server_mod._API_KEY = None
+
+ def test_correct_token_accepted(self, client):
+ _server_mod._API_KEY = "secret"
+ try:
+ r = client.get(
+ "/v1/collection/myidx/exists",
+ headers={"Authorization": "Bearer secret"},
+ )
+ assert r.status_code == 200
+ finally:
+ _server_mod._API_KEY = None
+
+ def test_health_exempt_from_auth(self, client):
+ _server_mod._API_KEY = "secret"
+ try:
+ r = client.get("/health")
+ assert r.status_code == 200
+ finally:
+ _server_mod._API_KEY = None
+
+
+# ---------------------------------------------------------------------------
+# Collection management
+# ---------------------------------------------------------------------------
+
+class TestCollectionManagement:
+ def test_collection_does_not_exist(self, client):
+ r = client.get("/v1/collection/noexist/exists")
+ assert r.status_code == 200
+ assert r.json()["exists"] is False
+
+ def test_open_creates_collection(self, client):
+ r = client.post(
+ "/v1/collection/open",
+ json={"index_name": "test_open", "backend": "local", "create": True, "dim": 8},
+ )
+ assert r.status_code == 200
+ body = r.json()
+ assert body["opened"] is True
+ assert body["created"] is True
+
+ def test_open_existing_collection(self, client):
+ client.post(
+ "/v1/collection/open",
+ json={"index_name": "test_reopen", "backend": "local", "create": True, "dim": 8},
+ )
+ r = client.post(
+ "/v1/collection/open",
+ json={"index_name": "test_reopen", "backend": "local", "create": False},
+ )
+ assert r.status_code == 200
+ assert r.json()["created"] is False
+
+ def test_open_nonexistent_without_create_returns_404(self, client):
+ r = client.post(
+ "/v1/collection/open",
+ json={"index_name": "missing", "backend": "local", "create": False},
+ )
+ assert r.status_code == 404
+
+ def test_create_collection_endpoint(self, client):
+ r = client.post(
+ "/v1/collection",
+ json={"index_name": "via_create", "backend": "local", "dim": 8},
+ )
+ assert r.status_code == 200
+ assert r.json()["created"] is True
+
+ def test_collection_exists_after_create(self, client):
+ client.post(
+ "/v1/collection",
+ json={"index_name": "check_exists", "backend": "local", "dim": 8},
+ )
+ r = client.get("/v1/collection/check_exists/exists")
+ assert r.json()["exists"] is True
+
+
+# ---------------------------------------------------------------------------
+# Upsert → search → fetch_vector full round-trip
+# ---------------------------------------------------------------------------
+
+class TestUpsertSearchFetch:
+ def _open(self, client, name="rt_test", backend="local", dim=8):
+ client.post(
+ "/v1/collection/open",
+ json={"index_name": name, "backend": backend, "create": True, "dim": dim},
+ )
+
+ def _upsert_points(self, client, name, n=3, dim=8):
+ points = []
+ for i in range(n):
+ pid = make_point_id(1, i)
+ points.append({
+ "point_id": pid,
+ "vector": torch.randn(5, dim),
+ "payload": {"doc_id": 1, "page_id": i},
+ })
+ body = _dumps(points)
+ return client.post(
+ f"/v1/upsert/{name}",
+ content=body,
+ headers={"Content-Type": "application/octet-stream"},
+ )
+
+ def test_upsert_returns_count(self, client):
+ self._open(client)
+ r = self._upsert_points(client, "rt_test", n=2)
+ assert r.status_code == 200
+ assert r.json()["upserted"] == 2
+
+ def test_point_exists_after_upsert(self, client):
+ self._open(client, name="pe_test")
+ self._upsert_points(client, "pe_test", n=1)
+ pid = make_point_id(1, 0)
+ r = client.get(f"/v1/point/pe_test/{pid}/exists")
+ assert r.json()["exists"] is True
+
+ def test_point_not_exists_before_upsert(self, client):
+ self._open(client, name="pne_test")
+ r = client.get(f"/v1/point/pne_test/9999999/exists")
+ assert r.json()["exists"] is False
+
+ def test_search_returns_hits(self, client):
+ self._open(client, name="srch_test")
+ self._upsert_points(client, "srch_test", n=3)
+
+ query_vec = torch.randn(4, 8)
+ wire = _dumps({"vectors": query_vec, "filter_metadata": None, "k": 2})
+ r = client.post(
+ "/v1/search/srch_test",
+ content=wire,
+ headers={"Content-Type": "application/octet-stream"},
+ )
+ assert r.status_code == 200
+ hits = _loads(r.content)
+ assert len(hits) <= 2
+ for h in hits:
+ assert "point_id" in h
+ assert "score" in h
+
+ def test_fetch_vector_returns_tensor(self, client):
+ self._open(client, name="fv_test")
+ self._upsert_points(client, "fv_test", n=1)
+ pid = make_point_id(1, 0)
+ r = client.get(f"/v1/vector/fv_test/{pid}")
+ assert r.status_code == 200
+ tensor = _loads(r.content)
+ assert tensor.shape[-1] == 8
+
+ def test_fetch_vector_missing_returns_404(self, client):
+ self._open(client, name="fv_missing")
+ r = client.get("/v1/vector/fv_missing/9999999")
+ assert r.status_code == 404
+
+ def test_search_on_unknown_collection_returns_404(self, client):
+ wire = _dumps({"vectors": torch.randn(2, 8), "filter_metadata": None, "k": 3})
+ r = client.post(
+ "/v1/search/no_such_col",
+ content=wire,
+ headers={"Content-Type": "application/octet-stream"},
+ )
+ assert r.status_code == 404
+
+
+# ---------------------------------------------------------------------------
+# Delete collection
+# ---------------------------------------------------------------------------
+
+class TestDeleteCollection:
+ def test_delete_removes_collection(self, client):
+ client.post(
+ "/v1/collection/open",
+ json={"index_name": "to_delete", "backend": "local", "create": True, "dim": 8},
+ )
+ r = client.delete("/v1/collection/to_delete")
+ assert r.status_code == 200
+ assert r.json()["deleted"] is True
+
+ r2 = client.get("/v1/collection/to_delete/exists")
+ assert r2.json()["exists"] is False
+
+ def test_delete_nonexistent_ok(self, client):
+ r = client.delete("/v1/collection/ghost_col")
+ assert r.status_code == 200
+
+
+# ---------------------------------------------------------------------------
+# Bookkeeping (server-side ColPali index metadata)
+# ---------------------------------------------------------------------------
+
+class TestBookkeeping:
+ def _blob(self):
+ return {
+ "index_config": {
+ "model_name": "vidore/colqwen-test",
+ "highest_doc_id": 4,
+ "description": "demo",
+ },
+ "embed_id_to_extra": {0: {"orig_size": (10, 20)}},
+ "doc_ids_to_file_names": {1: "a.pdf"},
+ "doc_id_to_metadata": {1: {"title": "A"}},
+ }
+
+ def test_get_missing_returns_404(self, client):
+ r = client.get("/v1/collection/bk_missing/bookkeeping")
+ assert r.status_code == 404
+
+ def test_put_then_get_roundtrip(self, client):
+ blob = self._blob()
+ r = client.put(
+ "/v1/collection/bk_rt/bookkeeping",
+ content=_dumps(blob),
+ headers={"Content-Type": "application/octet-stream"},
+ )
+ assert r.status_code == 200
+ assert r.json()["stored"] is True
+
+ r2 = client.get("/v1/collection/bk_rt/bookkeeping")
+ assert r2.status_code == 200
+ loaded = _loads(r2.content)
+ assert loaded["index_config"]["model_name"] == "vidore/colqwen-test"
+ assert loaded["doc_ids_to_file_names"][1] == "a.pdf"
+ assert loaded["doc_id_to_metadata"][1]["title"] == "A"
+
+ def test_put_overwrites(self, client):
+ client.put(
+ "/v1/collection/bk_ow/bookkeeping",
+ content=_dumps(self._blob()),
+ headers={"Content-Type": "application/octet-stream"},
+ )
+ updated = self._blob()
+ updated["index_config"]["description"] = "v2"
+ client.put(
+ "/v1/collection/bk_ow/bookkeeping",
+ content=_dumps(updated),
+ headers={"Content-Type": "application/octet-stream"},
+ )
+ r = client.get("/v1/collection/bk_ow/bookkeeping")
+ assert _loads(r.content)["index_config"]["description"] == "v2"
diff --git a/tests/test_vector_db_server_client.py b/tests/test_vector_db_server_client.py
new file mode 100644
index 0000000..832608c
--- /dev/null
+++ b/tests/test_vector_db_server_client.py
@@ -0,0 +1,293 @@
+"""Tests for VectorDBServerClient — HTTP interactions mocked via MagicMock."""
+
+import io
+from unittest.mock import MagicMock, patch
+import pytest
+import torch
+
+from foretrieval.vector_db_server.client import VectorDBServerClient, _dumps, _loads
+from foretrieval.vector_db_server.config import VectorDBServerConfig
+from foretrieval.vector_store.base import MultiVectorQuery, StoredPoint, make_point_id
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_config(**kwargs) -> VectorDBServerConfig:
+ defaults = {"url": "http://localhost:18000", "backend": "qdrant"}
+ return VectorDBServerConfig(**{**defaults, **kwargs})
+
+
+def _make_client(**cfg_kwargs) -> VectorDBServerClient:
+ return VectorDBServerClient(_make_config(**cfg_kwargs))
+
+
+def _make_mock_json_response(status_code: int, data: dict):
+ resp = MagicMock()
+ resp.status_code = status_code
+ resp.json.return_value = data
+ resp.text = str(data)
+ return resp
+
+
+def _make_mock_bytes_response(status_code: int, content: bytes):
+ resp = MagicMock()
+ resp.status_code = status_code
+ resp.content = content
+ resp.text = ""
+ resp.json.side_effect = Exception("not json")
+ return resp
+
+
+def _make_point(doc_id: int = 1, page_id: int = 0, dim: int = 4) -> StoredPoint:
+ pid = make_point_id(doc_id, page_id)
+ return StoredPoint(
+ point_id=pid,
+ vector=torch.randn(3, dim),
+ payload={"doc_id": doc_id, "page_id": page_id},
+ )
+
+
+# ---------------------------------------------------------------------------
+# Codec
+# ---------------------------------------------------------------------------
+
+class TestCodec:
+ def test_roundtrip_tensor(self):
+ t = torch.randn(5, 8)
+ assert torch.allclose(t, _loads(_dumps(t)))
+
+ def test_roundtrip_list_of_dicts(self):
+ data = [{"point_id": 1, "vector": torch.randn(3, 4), "payload": {"a": 1}}]
+ result = _loads(_dumps(data))
+ assert result[0]["point_id"] == 1
+ assert torch.allclose(result[0]["vector"], data[0]["vector"])
+
+
+# ---------------------------------------------------------------------------
+# Auth header
+# ---------------------------------------------------------------------------
+
+class TestAuthHeader:
+ def test_api_key_added_to_headers(self):
+ client = _make_client(api_key="mysecret")
+ assert "Authorization" in client._client.headers
+ assert client._client.headers["Authorization"] == "Bearer mysecret"
+
+ def test_no_api_key_no_auth_header(self):
+ client = _make_client()
+ assert "Authorization" not in client._client.headers
+
+
+# ---------------------------------------------------------------------------
+# Health check
+# ---------------------------------------------------------------------------
+
+class TestHealthCheck:
+ def test_health_ok(self):
+ client = _make_client()
+ client._client.get = MagicMock(
+ return_value=_make_mock_json_response(200, {"status": "ok"})
+ )
+ assert client.health_check() is True
+
+ def test_health_non_200(self):
+ client = _make_client()
+ client._client.get = MagicMock(
+ return_value=_make_mock_json_response(503, {})
+ )
+ assert client.health_check() is False
+
+ def test_health_connection_error(self):
+ import httpx
+ client = _make_client()
+ client._client.get = MagicMock(side_effect=httpx.ConnectError("refused"))
+ assert client.health_check() is False
+
+
+# ---------------------------------------------------------------------------
+# collection_exists
+# ---------------------------------------------------------------------------
+
+class TestCollectionExists:
+ def test_exists_true(self):
+ client = _make_client()
+ client._client.get = MagicMock(
+ return_value=_make_mock_json_response(200, {"exists": True, "backend": "qdrant"})
+ )
+ assert client.collection_exists("my_index") is True
+
+ def test_exists_false(self):
+ client = _make_client()
+ client._client.get = MagicMock(
+ return_value=_make_mock_json_response(200, {"exists": False, "backend": None})
+ )
+ assert client.collection_exists("my_index") is False
+
+
+# ---------------------------------------------------------------------------
+# open_collection
+# ---------------------------------------------------------------------------
+
+class TestOpenCollection:
+ def test_open_posts_correct_body(self):
+ client = _make_client()
+ expected = {"opened": True, "backend": "qdrant", "created": True}
+ client._client.post = MagicMock(
+ return_value=_make_mock_json_response(200, expected)
+ )
+ result = client.open_collection(
+ "idx", "qdrant", create=True, dim=128, storage_config=None
+ )
+ assert result == expected
+ call_args = client._client.post.call_args
+ assert "/v1/collection/open" in call_args.args[0]
+ body = call_args.kwargs["json"]
+ assert body["index_name"] == "idx"
+ assert body["create"] is True
+ assert body["dim"] == 128
+
+
+# ---------------------------------------------------------------------------
+# upsert
+# ---------------------------------------------------------------------------
+
+class TestUpsert:
+ def test_upsert_sends_octet_stream(self):
+ client = _make_client()
+ resp = _make_mock_json_response(200, {"upserted": 1})
+ client._client.post = MagicMock(return_value=resp)
+
+ point = _make_point()
+ client.upsert("idx", [point])
+
+ call_args = client._client.post.call_args
+ assert "upsert/idx" in call_args.args[0]
+ assert call_args.kwargs["headers"]["Content-Type"] == "application/octet-stream"
+ # Body is binary — decode and verify it contains the tensor
+ body_bytes = call_args.kwargs["content"]
+ decoded = _loads(body_bytes)
+ assert decoded[0]["point_id"] == point.point_id
+
+ def test_upsert_timeout_raises(self):
+ import httpx
+ client = _make_client()
+ client._client.post = MagicMock(side_effect=httpx.TimeoutException("timeout"))
+ with pytest.raises(TimeoutError):
+ client.upsert("idx", [_make_point()])
+
+ def test_upsert_connection_error_raises(self):
+ import httpx
+ client = _make_client()
+ client._client.post = MagicMock(side_effect=httpx.ConnectError("refused"))
+ with pytest.raises(ConnectionError):
+ client.upsert("idx", [_make_point()])
+
+
+# ---------------------------------------------------------------------------
+# point_exists
+# ---------------------------------------------------------------------------
+
+class TestPointExists:
+ def test_point_exists_true(self):
+ client = _make_client()
+ client._client.get = MagicMock(
+ return_value=_make_mock_json_response(200, {"exists": True})
+ )
+ assert client.point_exists("idx", 42) is True
+
+ def test_point_exists_false(self):
+ client = _make_client()
+ client._client.get = MagicMock(
+ return_value=_make_mock_json_response(200, {"exists": False})
+ )
+ assert client.point_exists("idx", 42) is False
+
+
+# ---------------------------------------------------------------------------
+# search
+# ---------------------------------------------------------------------------
+
+class TestSearch:
+ def test_search_sends_query_and_parses_hits(self):
+ client = _make_client()
+ hits_wire = [
+ {"point_id": 10000, "score": 9.5, "payload": {"doc_id": 1, "page_id": 0}}
+ ]
+ resp = _make_mock_bytes_response(200, _dumps(hits_wire))
+ client._client.post = MagicMock(return_value=resp)
+
+ query = MultiVectorQuery(vectors=torch.randn(5, 4))
+ hits = client.search("idx", query, k=3)
+
+ assert len(hits) == 1
+ assert hits[0].point_id == 10000
+ assert hits[0].score == pytest.approx(9.5)
+
+ # Verify request body contains k and vectors
+ body_bytes = client._client.post.call_args.kwargs["content"]
+ wire = _loads(body_bytes)
+ assert wire["k"] == 3
+ assert wire["filter_metadata"] is None
+
+
+# ---------------------------------------------------------------------------
+# fetch_vector
+# ---------------------------------------------------------------------------
+
+class TestFetchVector:
+ def test_fetch_returns_tensor(self):
+ client = _make_client()
+ tensor = torch.randn(3, 4)
+ resp = _make_mock_bytes_response(200, _dumps(tensor))
+ client._client.get = MagicMock(return_value=resp)
+
+ result = client.fetch_vector("idx", 12345)
+ assert result is not None
+ assert torch.allclose(result, tensor)
+
+ def test_fetch_returns_none_on_404(self):
+ client = _make_client()
+ resp = MagicMock()
+ resp.status_code = 404
+ client._client.get = MagicMock(return_value=resp)
+
+ result = client.fetch_vector("idx", 99999)
+ assert result is None
+
+
+# ---------------------------------------------------------------------------
+# Bookkeeping
+# ---------------------------------------------------------------------------
+
+class TestBookkeeping:
+ def test_put_bookkeeping_sends_octet_stream(self):
+ client = _make_client()
+ resp = _make_mock_json_response(200, {"stored": True})
+ client._client.put = MagicMock(return_value=resp)
+
+ blob = {"index_config": {"model_name": "m"}, "doc_id_to_metadata": {1: {"t": "x"}}}
+ client.put_bookkeeping("idx", blob)
+
+ call = client._client.put.call_args
+ assert call.args[0].endswith("/v1/collection/idx/bookkeeping")
+ body = call.kwargs["content"]
+ assert _loads(body)["index_config"]["model_name"] == "m"
+
+ def test_get_bookkeeping_returns_blob(self):
+ client = _make_client()
+ blob = {"index_config": {"model_name": "m"}}
+ resp = _make_mock_bytes_response(200, _dumps(blob))
+ client._client.get = MagicMock(return_value=resp)
+
+ out = client.get_bookkeeping("idx")
+ assert out["index_config"]["model_name"] == "m"
+
+ def test_get_bookkeeping_returns_none_on_404(self):
+ client = _make_client()
+ resp = MagicMock()
+ resp.status_code = 404
+ client._client.get = MagicMock(return_value=resp)
+
+ assert client.get_bookkeeping("idx") is None
diff --git a/tests/test_vector_db_server_config.py b/tests/test_vector_db_server_config.py
new file mode 100644
index 0000000..ce0aad2
--- /dev/null
+++ b/tests/test_vector_db_server_config.py
@@ -0,0 +1,104 @@
+"""Tests for VectorDBServerConfig — pydantic validators and helpers."""
+
+import pytest
+from pydantic import ValidationError
+
+from foretrieval.vector_db_server.config import VectorDBServerConfig
+
+
+def _make_cfg(**kwargs) -> VectorDBServerConfig:
+ defaults = {"url": "http://localhost:18000", "backend": "qdrant"}
+ return VectorDBServerConfig(**{**defaults, **kwargs})
+
+
+class TestURLValidation:
+ def test_trailing_slash_stripped(self):
+ cfg = _make_cfg(url="http://db-host:18000/")
+ assert cfg.url == "http://db-host:18000"
+
+ def test_multiple_trailing_slashes_stripped(self):
+ cfg = _make_cfg(url="http://db-host:18000///")
+ assert cfg.url == "http://db-host:18000"
+
+ def test_url_preserved_without_trailing_slash(self):
+ cfg = _make_cfg(url="http://db-host:18000")
+ assert cfg.url == "http://db-host:18000"
+
+
+class TestBackendValidation:
+ def test_valid_backends(self):
+ for b in ("local", "qdrant", "milvus"):
+ cfg = _make_cfg(backend=b)
+ assert cfg.backend == b
+
+ def test_backend_normalised_to_lower(self):
+ cfg = _make_cfg(backend="Qdrant")
+ assert cfg.backend == "qdrant"
+
+ def test_invalid_backend_raises(self):
+ with pytest.raises(ValidationError, match="not a valid server-side backend"):
+ _make_cfg(backend="redis")
+
+
+class TestPortValidation:
+ def test_default_port(self):
+ cfg = _make_cfg()
+ assert cfg.port == 18000
+
+ def test_custom_port(self):
+ cfg = _make_cfg(port=9999)
+ assert cfg.port == 9999
+
+ def test_zero_port_raises(self):
+ with pytest.raises(ValidationError):
+ _make_cfg(port=0)
+
+ def test_out_of_range_port_raises(self):
+ with pytest.raises(ValidationError):
+ _make_cfg(port=99999)
+
+
+class TestAutoDeployValidation:
+ def test_auto_deploy_requires_ssh_host(self):
+ with pytest.raises(ValidationError, match="ssh_host is required"):
+ _make_cfg(auto_deploy=True, ssh_host=None)
+
+ def test_auto_deploy_with_ssh_host_ok(self):
+ cfg = _make_cfg(auto_deploy=True, ssh_host="gpu-server")
+ assert cfg.auto_deploy is True
+ assert cfg.ssh_host == "gpu-server"
+
+ def test_auto_deploy_false_no_ssh_needed(self):
+ cfg = _make_cfg(auto_deploy=False)
+ assert cfg.auto_deploy is False
+
+
+class TestDefaults:
+ def test_api_key_default_none(self):
+ cfg = _make_cfg()
+ assert cfg.api_key is None
+
+ def test_verify_ssl_default_true(self):
+ cfg = _make_cfg()
+ assert cfg.verify_ssl is True
+
+ def test_request_timeout_default(self):
+ cfg = _make_cfg()
+ assert cfg.request_timeout == 120
+
+ def test_data_dir_default(self):
+ cfg = _make_cfg()
+ assert cfg.data_dir == "/var/lib/foretrieval_db"
+
+
+class TestFromDict:
+ def test_from_dict_basic(self):
+ cfg = VectorDBServerConfig.from_dict(
+ {"url": "http://db-server:18000", "backend": "qdrant", "api_key": "tok"}
+ )
+ assert cfg.url == "http://db-server:18000"
+ assert cfg.api_key == "tok"
+
+ def test_from_dict_missing_url_raises(self):
+ with pytest.raises((ValidationError, TypeError)):
+ VectorDBServerConfig.from_dict({"backend": "qdrant"})
diff --git a/tests/test_vector_db_server_manager.py b/tests/test_vector_db_server_manager.py
new file mode 100644
index 0000000..ef3c3c3
--- /dev/null
+++ b/tests/test_vector_db_server_manager.py
@@ -0,0 +1,446 @@
+"""Tests for VectorDBServerManager — SSH/Docker operations mocked via MagicMock."""
+
+import json
+from unittest.mock import MagicMock, patch
+import pytest
+
+from foretrieval.vector_db_server.config import VectorDBServerConfig
+from foretrieval.vector_db_server.manager import VectorDBServerManager, _CONTAINER_NAME
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_manager(**cfg_kwargs) -> VectorDBServerManager:
+ defaults = {
+ "url": "http://localhost:18000",
+ "backend": "qdrant",
+ "ssh_host": "gpu-server",
+ "auto_deploy": True,
+ }
+ cfg = VectorDBServerConfig(**{**defaults, **cfg_kwargs})
+ mgr = VectorDBServerManager(cfg)
+ # Stub the remote-home resolution so tests that mock _run_remote
+ # don't accidentally trigger a real SSH connection through
+ # _remote_home() -> sftp.normalize('.').
+ mgr._cached_home = "/home/testuser"
+ # Stub UID/GID so _build_docker_run_cmd tests don't need SSH.
+ # Tests that specifically exercise the --user flag override this.
+ mgr._resolve_remote_uid_gid = MagicMock(return_value=(None, None))
+ return mgr
+
+
+def _fake_run_factory(responses: dict):
+ """Return a _run_remote mock whose stdout depends on substrings in the command."""
+ def _fake_run(cmd: str, on_line=None) -> tuple[str, str]:
+ for substring, output in responses.items():
+ if substring in cmd:
+ return output, ""
+ return "", ""
+ return _fake_run
+
+
+# ---------------------------------------------------------------------------
+# Constructor
+# ---------------------------------------------------------------------------
+
+class TestManagerConstructor:
+ def test_requires_ssh_host(self):
+ cfg = VectorDBServerConfig(
+ url="http://localhost:18000", backend="qdrant", ssh_host=None
+ )
+ with pytest.raises(ValueError, match="ssh_host"):
+ VectorDBServerManager(cfg)
+
+ def test_paramiko_import_error(self):
+ mgr = _make_manager()
+ mgr._run_remote = MagicMock()
+ with patch.dict("sys.modules", {"paramiko": None}):
+ with pytest.raises(ImportError, match="paramiko"):
+ mgr.ensure_deployed()
+
+
+# ---------------------------------------------------------------------------
+# ensure_deployed — deploy from scratch
+# ---------------------------------------------------------------------------
+
+class TestDeployFromScratch:
+ def test_deploys_when_no_metadata(self, tmp_path):
+ mgr = _make_manager()
+
+ # _run_remote returns __MISSING__ for metadata cat, "true" for docker inspect
+ calls = []
+ def fake_run(cmd: str, on_line=None):
+ calls.append(cmd)
+ if "cat" in cmd and "db_deployment" in cmd:
+ return "__MISSING__", ""
+ if "docker inspect" in cmd:
+ return "false", ""
+ if "nvidia-smi" in cmd:
+ return "2\n", ""
+ return "", ""
+
+ mgr._run_remote = MagicMock(side_effect=fake_run)
+ # Mock _upload_build_context to avoid real filesystem + SSH
+ mgr._upload_build_context = MagicMock()
+
+ # Stub paramiko import check
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.ensure_deployed()
+
+ # Should have called upload + docker build + docker run
+ all_cmds = " ".join(calls)
+ assert "docker build" in all_cmds
+ assert "docker run" in all_cmds
+
+ def test_no_op_when_container_running(self):
+ mgr = _make_manager()
+
+ meta = json.dumps({"deployed_at": "2026-01-01T00:00:00+00:00"})
+
+ def fake_run(cmd: str, on_line=None):
+ if "cat" in cmd and "db_deployment" in cmd:
+ return meta, ""
+ if "docker inspect" in cmd:
+ return "true", ""
+ return "", ""
+
+ mgr._run_remote = MagicMock(side_effect=fake_run)
+ mgr._upload_build_context = MagicMock()
+
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.ensure_deployed()
+
+ # _upload_build_context should NOT have been called
+ mgr._upload_build_context.assert_not_called()
+
+ def test_redeploys_when_container_stopped(self):
+ mgr = _make_manager()
+
+ meta = json.dumps({"deployed_at": "2026-01-01T00:00:00+00:00"})
+ calls = []
+
+ def fake_run(cmd: str, on_line=None):
+ calls.append(cmd)
+ if "cat" in cmd and "db_deployment" in cmd:
+ return meta, ""
+ if "docker inspect" in cmd:
+ return "false", ""
+ return "", ""
+
+ mgr._run_remote = MagicMock(side_effect=fake_run)
+ mgr._upload_build_context = MagicMock()
+
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.ensure_deployed()
+
+ all_cmds = " ".join(calls)
+ assert "docker build" in all_cmds
+
+
+# ---------------------------------------------------------------------------
+# _build_docker_run_cmd
+# ---------------------------------------------------------------------------
+
+class TestBuildDockerRunCmd:
+ def test_port_in_command(self):
+ mgr = _make_manager(port=18000)
+ cmd = mgr._build_docker_run_cmd()
+ assert "-p 18000:18000" in cmd
+
+ def test_data_dir_mount(self):
+ mgr = _make_manager(data_dir="/mnt/data")
+ cmd = mgr._build_docker_run_cmd()
+ assert "/mnt/data:/data" in cmd
+
+ def test_api_key_in_env(self):
+ mgr = _make_manager(api_key="mysecret")
+ cmd = mgr._build_docker_run_cmd()
+ assert "FOR_DB_API_KEY=mysecret" in cmd
+
+ def test_no_api_key_not_in_cmd(self):
+ mgr = _make_manager()
+ cmd = mgr._build_docker_run_cmd()
+ assert "FOR_DB_API_KEY" not in cmd
+
+ def test_container_name_in_cmd(self):
+ mgr = _make_manager()
+ cmd = mgr._build_docker_run_cmd()
+ assert f"--name {_CONTAINER_NAME}" in cmd
+
+ def test_user_flag_present_when_uid_resolved(self):
+ """When UID/GID can be resolved, --user uid:gid must appear."""
+ mgr = _make_manager()
+ mgr._resolve_remote_uid_gid = MagicMock(return_value=(1000, 1001))
+ cmd = mgr._build_docker_run_cmd()
+ assert "--user 1000:1001" in cmd
+
+ def test_user_flag_absent_when_uid_unknown(self):
+ """When resolution fails (returns None, None), no --user flag."""
+ mgr = _make_manager()
+ # _resolve_remote_uid_gid already returns (None, None) from _make_manager
+ cmd = mgr._build_docker_run_cmd()
+ assert "--user" not in cmd
+
+
+# ---------------------------------------------------------------------------
+# stop
+# ---------------------------------------------------------------------------
+
+class TestStop:
+ def test_stop_sends_correct_commands(self):
+ mgr = _make_manager()
+ calls = []
+ mgr._run_remote = MagicMock(side_effect=lambda cmd, on_line=None: (calls.append(cmd), ("", ""))[1])
+ mgr.stop()
+ all_cmds = " ".join(calls)
+ assert "docker stop" in all_cmds
+ assert "docker rm" in all_cmds
+ assert "db_deployment" in all_cmds
+
+
+# ---------------------------------------------------------------------------
+# redeploy + on_line streaming
+# ---------------------------------------------------------------------------
+
+class TestRedeploy:
+ def test_redeploy_calls_deploy(self):
+ mgr = _make_manager()
+ mgr._deploy = MagicMock()
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.redeploy()
+ mgr._deploy.assert_called_once()
+
+ def test_redeploy_forwards_on_line_callback(self):
+ mgr = _make_manager()
+ mgr._deploy = MagicMock()
+ cb = MagicMock()
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.redeploy(on_line=cb)
+ mgr._deploy.assert_called_once()
+ # The callback was threaded through
+ assert mgr._deploy.call_args.kwargs.get("on_line") is cb
+
+ def test_redeploy_unconditionally_redeploys_running_container(self):
+ """Unlike ensure_deployed, redeploy() always rebuilds."""
+ mgr = _make_manager()
+
+ # _is_container_running returns True; ensure_deployed would no-op,
+ # but redeploy must still trigger _deploy.
+ mgr._is_container_running = MagicMock(return_value=True)
+ mgr._read_remote_metadata = MagicMock(return_value={"deployed_at": "x"})
+ mgr._deploy = MagicMock()
+ with patch.dict("sys.modules", {"paramiko": MagicMock()}):
+ mgr.redeploy()
+ mgr._deploy.assert_called_once()
+
+
+class TestOnLineStreaming:
+ def test_run_remote_streams_lines_when_callback_given(self):
+ """When on_line is provided, stdout is read line-by-line."""
+ mgr = _make_manager()
+
+ # Build a fake SSH client whose exec_command returns three lines
+ class _FakeChannel:
+ def recv_exit_status(self):
+ return 0
+
+ class _FakeStdout:
+ def __init__(self, lines):
+ self._lines = list(lines)
+ self.channel = _FakeChannel()
+
+ def readline(self):
+ if self._lines:
+ return self._lines.pop(0)
+ return ""
+
+ def read(self):
+ return b""
+
+ fake_stdout = _FakeStdout(["Step 1/3\n", "Step 2/3\n", "Done\n"])
+ fake_stderr = MagicMock()
+ fake_stderr.read.return_value = b""
+ fake_ssh = MagicMock()
+ fake_ssh.exec_command.return_value = (None, fake_stdout, fake_stderr)
+ mgr._get_ssh = MagicMock(return_value=fake_ssh)
+
+ captured = []
+ stdout, stderr = mgr._run_remote("docker build .", on_line=captured.append)
+
+ assert captured == ["Step 1/3", "Step 2/3", "Done"]
+ assert "Step 1/3" in stdout
+ assert stderr == ""
+
+ def test_run_remote_buffered_when_no_callback(self):
+ """Default behaviour (no on_line) still reads the whole stdout at once."""
+ mgr = _make_manager()
+
+ class _FakeChannel:
+ def recv_exit_status(self):
+ return 0
+
+ class _FakeStdout:
+ channel = _FakeChannel()
+ def read(self):
+ return b"full output\n"
+
+ fake_stderr = MagicMock()
+ fake_stderr.read.return_value = b""
+ fake_ssh = MagicMock()
+ fake_ssh.exec_command.return_value = (None, _FakeStdout(), fake_stderr)
+ mgr._get_ssh = MagicMock(return_value=fake_ssh)
+
+ stdout, _ = mgr._run_remote("echo hi")
+ assert "full output" in stdout
+
+
+# ---------------------------------------------------------------------------
+# _get_ssh uses ssh_utils.open_ssh_client
+# ---------------------------------------------------------------------------
+
+class TestGetSshUsesSshUtils:
+ def test_get_ssh_delegates_to_open_ssh_client(self):
+ """_get_ssh must call foretrieval.ssh_utils.open_ssh_client.
+
+ This is the regression test for the bug where paramiko was
+ called with a raw alias string and DNS-failed for entries that
+ only existed in ~/.ssh/config.
+ """
+ mgr = _make_manager(ssh_host="pf01", ssh_user="alice", ssh_key_path="/k")
+ fake_client = MagicMock()
+ with patch("foretrieval.ssh_utils.open_ssh_client",
+ return_value=fake_client) as mock_open:
+ result = mgr._get_ssh()
+ mock_open.assert_called_once_with(
+ ssh_host="pf01",
+ ssh_user="alice",
+ ssh_key_path="/k",
+ )
+ assert result is fake_client
+
+ def test_get_ssh_is_cached(self):
+ mgr = _make_manager()
+ fake_client = MagicMock()
+ with patch("foretrieval.ssh_utils.open_ssh_client",
+ return_value=fake_client) as mock_open:
+ r1 = mgr._get_ssh()
+ r2 = mgr._get_ssh()
+ assert r1 is r2
+ assert mock_open.call_count == 1
+
+
+# ---------------------------------------------------------------------------
+# Remote home / metadata path resolution
+# ---------------------------------------------------------------------------
+
+class TestRemoteHome:
+ def test_remote_home_uses_sftp_normalize(self):
+ """_remote_home returns the path produced by sftp.normalize('.')."""
+ mgr = _make_manager()
+ mgr._cached_home = None # reset the stub from _make_manager
+
+ fake_sftp = MagicMock()
+ fake_sftp.normalize.return_value = "/home/alice"
+ fake_ssh = MagicMock()
+ fake_ssh.open_sftp.return_value = fake_sftp
+ mgr._get_ssh = MagicMock(return_value=fake_ssh)
+
+ assert mgr._remote_home() == "/home/alice"
+ fake_sftp.normalize.assert_called_once_with(".")
+ fake_sftp.close.assert_called_once()
+
+ def test_remote_home_caches_result(self):
+ mgr = _make_manager()
+ mgr._cached_home = None
+
+ fake_sftp = MagicMock()
+ fake_sftp.normalize.return_value = "/home/alice"
+ fake_ssh = MagicMock()
+ fake_ssh.open_sftp.return_value = fake_sftp
+ mgr._get_ssh = MagicMock(return_value=fake_ssh)
+
+ a = mgr._remote_home()
+ b = mgr._remote_home()
+ assert a == b == "/home/alice"
+ fake_ssh.open_sftp.assert_called_once()
+
+
+class TestMetadataPathAbsolute:
+ def test_metadata_path_is_absolute(self):
+ mgr = _make_manager() # _cached_home = "/home/testuser"
+ path = mgr._metadata_path()
+ assert path.startswith("/home/testuser/")
+ assert "~" not in path
+ assert path.endswith("db_deployment.json")
+
+ def test_stop_uses_absolute_metadata_path(self):
+ """Regression: stop() must rm an absolute path, not '~/...'."""
+ mgr = _make_manager()
+ calls = []
+ mgr._run_remote = MagicMock(
+ side_effect=lambda cmd, on_line=None: (calls.append(cmd), ("", ""))[1]
+ )
+ mgr.stop()
+ rm_calls = [c for c in calls if c.startswith("rm -f")]
+ assert rm_calls, "stop() did not call rm -f"
+ # No tilde in the path passed to rm
+ assert all("~" not in c for c in rm_calls), rm_calls
+ assert any("/home/testuser/" in c for c in rm_calls), rm_calls
+
+
+class TestUploadBuildContextAbsolutePath:
+ def test_sftp_put_uses_absolute_remote_path(self, tmp_path, monkeypatch):
+ """Regression for the SFTP ENOENT bug.
+
+ The previous implementation passed
+ f"{'~/foretrieval_db_build'.replace('~', '')}/build_context.tar.gz"
+ which resolved to '/foretrieval_db_build/...' (root-relative).
+ The new code must pass an absolute path rooted at the remote home.
+ """
+ from foretrieval.vector_db_server import manager as mod
+
+ mgr = _make_manager() # _cached_home = "/home/testuser"
+
+ # Mock _run_remote (mkdir, tar extract)
+ runs: list[str] = []
+ mgr._run_remote = MagicMock(
+ side_effect=lambda cmd, on_line=None: (runs.append(cmd), ("", ""))[1]
+ )
+
+ # Mock SSH + SFTP
+ sftp_calls: list[tuple] = []
+
+ class _FakeSftp:
+ def put(self, local, remote):
+ sftp_calls.append(("put", local, remote))
+ def close(self):
+ sftp_calls.append(("close",))
+
+ fake_sftp = _FakeSftp()
+ fake_ssh = MagicMock()
+ fake_ssh.open_sftp.return_value = fake_sftp
+ mgr._get_ssh = MagicMock(return_value=fake_ssh)
+
+ # Avoid real disk + real tar — short-circuit _upload_build_context's
+ # source-locating step by pointing it at the actual repo (it exists).
+ # tarfile is small so let it run.
+ mgr._upload_build_context()
+
+ # Verify SFTP put destination
+ put_calls = [c for c in sftp_calls if c[0] == "put"]
+ assert len(put_calls) == 1
+ remote_dest = put_calls[0][2]
+ assert remote_dest == "/home/testuser/foretrieval_db_build/build_context.tar.gz", remote_dest
+
+ # Verify the mkdir uses the same absolute path
+ mkdir_calls = [c for c in runs if c.startswith("mkdir -p")]
+ assert any("/home/testuser/foretrieval_db_build" in c for c in mkdir_calls), mkdir_calls
+
+ # Verify tar extraction targets the same absolute path
+ tar_calls = [c for c in runs if "tar -xzf" in c]
+ assert any("/home/testuser/foretrieval_db_build" in c for c in tar_calls), tar_calls
+
+ # _remote_build_dir stashed for _deploy() to reuse
+ assert mgr._remote_build_dir == "/home/testuser/foretrieval_db_build"
diff --git a/tests/test_vector_store_factory.py b/tests/test_vector_store_factory.py
new file mode 100644
index 0000000..f983e6a
--- /dev/null
+++ b/tests/test_vector_store_factory.py
@@ -0,0 +1,143 @@
+"""Tests for make_vector_store() factory."""
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+import pytest
+
+from foretrieval.vector_store.factory import make_vector_store, BACKEND_REGISTRY
+from foretrieval.vector_store.local import LocalVectorStore
+from foretrieval.vector_store.qdrant import QdrantVectorStore
+from foretrieval.vector_store.milvus import MilvusVectorStore
+from foretrieval.vector_store.remote import RemoteVectorStore
+
+
+class TestFactory:
+ def test_local_backend_returns_local_store(self):
+ vs = make_vector_store("local")
+ assert isinstance(vs, LocalVectorStore)
+
+ def test_qdrant_backend_returns_qdrant_store(self):
+ vs = make_vector_store("qdrant")
+ assert isinstance(vs, QdrantVectorStore)
+
+ def test_milvus_backend_returns_milvus_store(self):
+ vs = make_vector_store("milvus")
+ assert isinstance(vs, MilvusVectorStore)
+
+ def test_unknown_backend_raises_value_error(self):
+ with pytest.raises(ValueError, match="Unknown storage backend"):
+ make_vector_store("nonexistent")
+
+ def test_error_message_lists_supported_backends(self):
+ with pytest.raises(ValueError, match="local"):
+ make_vector_store("bad")
+
+ def test_none_backend_defaults_to_local(self):
+ vs = make_vector_store(None)
+ assert isinstance(vs, LocalVectorStore)
+
+ def test_whitespace_and_uppercase_handled(self):
+ vs = make_vector_store(" LOCAL ")
+ assert isinstance(vs, LocalVectorStore)
+
+ def test_milvus_candidate_limit_passed_through(self):
+ vs = make_vector_store("milvus", {"candidate_limit": 128})
+ assert isinstance(vs, MilvusVectorStore)
+ assert vs._candidate_limit == 128
+
+ def test_backend_registry_contains_all_four(self):
+ assert "local" in BACKEND_REGISTRY
+ assert "qdrant" in BACKEND_REGISTRY
+ assert "milvus" in BACKEND_REGISTRY
+ assert "remote" in BACKEND_REGISTRY
+
+ def test_each_call_returns_new_instance(self):
+ a = make_vector_store("local")
+ b = make_vector_store("local")
+ assert a is not b
+
+
+class TestRemoteFactory:
+ """Test make_vector_store("remote", ...) factory path."""
+
+ def _make_remote(self, extra=None, auto_deploy=False):
+ cfg = {
+ "url": "http://localhost:18000",
+ "backend": "qdrant",
+ "auto_deploy": auto_deploy,
+ }
+ if extra:
+ cfg.update(extra)
+ # Patch VectorDBServerClient inside the module where it is imported
+ with patch(
+ "foretrieval.vector_db_server.client.VectorDBServerClient"
+ ):
+ with patch(
+ "foretrieval.vector_store.factory._make_remote_vector_store",
+ ) as mock_factory:
+ mock_client_inst = MagicMock()
+ remote_store = RemoteVectorStore(mock_client_inst, backend=cfg.get("backend", "qdrant"))
+ mock_factory.return_value = remote_store
+ vs = make_vector_store("remote", cfg)
+ return vs, mock_client_inst
+
+ def test_remote_backend_returns_remote_store(self):
+ vs, _ = self._make_remote()
+ assert isinstance(vs, RemoteVectorStore)
+
+ def test_remote_store_has_correct_server_backend(self):
+ vs, _ = self._make_remote()
+ assert vs.server_backend == "qdrant"
+
+ def test_remote_store_milvus_backend(self):
+ vs, _ = self._make_remote(extra={"backend": "milvus"})
+ assert vs.server_backend == "milvus"
+
+ def test_remote_missing_url_raises(self):
+ with pytest.raises(ValueError, match="url"):
+ make_vector_store("remote", {"backend": "qdrant"})
+
+ def test_auto_deploy_calls_manager_and_health_check(self):
+ with (
+ patch("foretrieval.vector_db_server.client.VectorDBServerClient") as MockClient,
+ patch("foretrieval.vector_db_server.manager.VectorDBServerManager") as MockMgr,
+ patch("foretrieval.vector_store.factory._wait_for_health") as mock_wait,
+ ):
+ mock_client_inst = MagicMock()
+ MockClient.return_value = mock_client_inst
+ mock_mgr_inst = MagicMock()
+ MockMgr.return_value = mock_mgr_inst
+
+ vs = make_vector_store(
+ "remote",
+ {
+ "url": "http://gpu-server:18000",
+ "backend": "qdrant",
+ "auto_deploy": True,
+ "ssh_host": "gpu-server",
+ },
+ )
+ mock_mgr_inst.ensure_deployed.assert_called_once()
+ mock_wait.assert_called_once()
+ assert isinstance(vs, RemoteVectorStore)
+
+ def test_candidate_limit_forwarded_as_server_storage_config(self):
+ mock_client_inst = MagicMock()
+ remote_store = RemoteVectorStore(
+ mock_client_inst, backend="milvus",
+ storage_config={"candidate_limit": 128}
+ )
+ with patch(
+ "foretrieval.vector_store.factory._make_remote_vector_store",
+ return_value=remote_store,
+ ):
+ vs = make_vector_store(
+ "remote",
+ {
+ "url": "http://localhost:18000",
+ "backend": "milvus",
+ "candidate_limit": 128,
+ },
+ )
+ assert vs._storage_config == {"candidate_limit": 128}
+
diff --git a/tests/test_vector_store_local.py b/tests/test_vector_store_local.py
new file mode 100644
index 0000000..0142c38
--- /dev/null
+++ b/tests/test_vector_store_local.py
@@ -0,0 +1,258 @@
+"""Tests for LocalVectorStore.
+
+All tests run without GPU — processor.score() is mocked.
+"""
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+
+from foretrieval.vector_store.base import (
+ MultiVectorQuery,
+ StoredPoint,
+ make_point_id,
+)
+from foretrieval.vector_store.local import LocalVectorStore
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_store(processor=None) -> LocalVectorStore:
+ store = LocalVectorStore()
+ if processor is not None:
+ store.set_processor(processor)
+ return store
+
+
+def _dummy_embedding(n_tokens: int = 4, dim: int = 8) -> torch.Tensor:
+ return torch.rand(n_tokens, dim)
+
+
+def _mock_processor(scores: list[float]):
+ """Return a processor mock whose score() returns the given scores as a 2-D array."""
+ proc = MagicMock()
+ proc.score.return_value = torch.tensor([scores])
+ return proc
+
+
+def _point(doc_id: int, page_id: int, vec: torch.Tensor | None = None) -> StoredPoint:
+ if vec is None:
+ vec = _dummy_embedding()
+ pid = make_point_id(doc_id, page_id)
+ return StoredPoint(
+ point_id=pid,
+ vector=vec,
+ payload={"doc_id": doc_id, "page_id": page_id, "chunk_id": None, "metadata": {}},
+ )
+
+
+# ---------------------------------------------------------------------------
+# Lifecycle
+# ---------------------------------------------------------------------------
+
+class TestLifecycle:
+ def test_open_does_not_raise(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True, dim=8)
+
+ def test_collection_not_exists_before_upsert(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True, dim=8)
+ assert not store.collection_exists()
+
+ def test_collection_exists_after_export(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True, dim=8)
+ store.upsert([_point(0, 1)])
+ index_dir = tmp_path / "idx"
+ index_dir.mkdir(parents=True, exist_ok=True)
+ store.export_sidecar(index_dir)
+ assert store.collection_exists()
+
+
+# ---------------------------------------------------------------------------
+# Upsert and point_exists
+# ---------------------------------------------------------------------------
+
+class TestUpsert:
+ def test_upsert_single_point(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ sp = _point(0, 1)
+ store.upsert([sp])
+ assert store.point_exists(sp.point_id)
+
+ def test_upsert_multiple_points(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ pts = [_point(i, 1) for i in range(5)]
+ store.upsert(pts)
+ for p in pts:
+ assert store.point_exists(p.point_id)
+
+ def test_point_not_exists_before_upsert(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ assert not store.point_exists(make_point_id(0, 1))
+
+ def test_upsert_idempotent_last_write_wins(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ pid = make_point_id(0, 1)
+ v1 = torch.zeros(4, 8)
+ v2 = torch.ones(4, 8)
+ sp1 = StoredPoint(point_id=pid, vector=v1, payload={"doc_id": 0, "page_id": 1, "chunk_id": None, "metadata": {}})
+ sp2 = StoredPoint(point_id=pid, vector=v2, payload={"doc_id": 0, "page_id": 1, "chunk_id": None, "metadata": {}})
+ store.upsert([sp1])
+ store.upsert([sp2])
+ fetched = store.fetch_vector(pid)
+ assert fetched is not None
+ assert torch.allclose(fetched, v2)
+
+
+# ---------------------------------------------------------------------------
+# Search
+# ---------------------------------------------------------------------------
+
+class TestSearch:
+ def test_search_returns_k_results(self, tmp_path):
+ scores = [0.5, 0.8, 0.3]
+ proc = _mock_processor(scores)
+ store = _make_store(proc)
+ store.open("idx", tmp_path, create=True)
+ for i in range(3):
+ store.upsert([_point(i, 1)])
+ q = MultiVectorQuery(vectors=_dummy_embedding(2, 8))
+ results = store.search(q, k=2)
+ assert len(results) == 2
+
+ def test_search_sorted_descending_score(self, tmp_path):
+ scores = [0.3, 0.9, 0.1]
+ proc = _mock_processor(scores)
+ store = _make_store(proc)
+ store.open("idx", tmp_path, create=True)
+ for i in range(3):
+ store.upsert([_point(i, 1)])
+ q = MultiVectorQuery(vectors=_dummy_embedding(2, 8))
+ results = store.search(q, k=3)
+ assert results[0].score >= results[1].score >= results[2].score
+
+ def test_search_returns_correct_point_ids(self, tmp_path):
+ scores = [0.1, 0.9, 0.5]
+ proc = _mock_processor(scores)
+ store = _make_store(proc)
+ store.open("idx", tmp_path, create=True)
+ pts = [_point(i, 1) for i in range(3)]
+ for p in pts:
+ store.upsert([p])
+ q = MultiVectorQuery(vectors=_dummy_embedding(2, 8))
+ results = store.search(q, k=1)
+ # Best score is index 1 (score 0.9 → doc_id=1)
+ assert results[0].payload["doc_id"] == 1
+
+ def test_search_raises_without_processor(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ store.upsert([_point(0, 1)])
+ q = MultiVectorQuery(vectors=_dummy_embedding(2, 8))
+ with pytest.raises(RuntimeError, match="set_processor"):
+ store.search(q, k=1)
+
+ def test_search_empty_store_returns_empty(self, tmp_path):
+ proc = MagicMock()
+ store = _make_store(proc)
+ store.open("idx", tmp_path, create=True)
+ q = MultiVectorQuery(vectors=_dummy_embedding(2, 8))
+ assert store.search(q, k=5) == []
+
+ def test_search_with_metadata_filter(self, tmp_path):
+ """Filter should restrict search to matching docs only."""
+ # 3 docs: doc 0 has language=en, doc 1 has language=fr, doc 2 has language=en
+ # We filter on language=fr → only doc 1 should be in the result pool.
+ # The mock processor returns a single score for one embedding (the filtered one).
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ store.set_doc_id_to_metadata({
+ 0: {"language": "en"},
+ 1: {"language": "fr"},
+ 2: {"language": "en"},
+ })
+ for i in range(3):
+ store.upsert([_point(i, 1)])
+
+ # After metadata filter only doc 1 is left → processor.score gets 1 embedding
+ proc_filtered = MagicMock()
+ proc_filtered.score.return_value = torch.tensor([[0.7]])
+ store.set_processor(proc_filtered)
+
+ q = MultiVectorQuery(
+ vectors=_dummy_embedding(2, 8),
+ filter_metadata={"language": "fr"},
+ )
+ results = store.search(q, k=5)
+ assert len(results) == 1
+ assert results[0].payload["doc_id"] == 1
+
+
+# ---------------------------------------------------------------------------
+# fetch_vector
+# ---------------------------------------------------------------------------
+
+class TestFetchVector:
+ def test_fetch_returns_tensor(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ vec = _dummy_embedding(4, 8)
+ sp = _point(0, 1, vec)
+ store.upsert([sp])
+ fetched = store.fetch_vector(sp.point_id)
+ assert fetched is not None
+ assert fetched.shape == (4, 8)
+ assert torch.allclose(fetched, vec)
+
+ def test_fetch_unknown_id_returns_none(self, tmp_path):
+ store = LocalVectorStore()
+ store.open("idx", tmp_path, create=True)
+ assert store.fetch_vector(make_point_id(99, 99)) is None
+
+
+# ---------------------------------------------------------------------------
+# Persistence round-trip
+# ---------------------------------------------------------------------------
+
+class TestPersistence:
+ def test_export_load_round_trip(self, tmp_path):
+ store1 = LocalVectorStore()
+ store1.open("idx", tmp_path, create=True)
+ pts = [_point(i, 1) for i in range(3)]
+ for p in pts:
+ store1.upsert([p])
+ index_path = tmp_path / "idx"
+ index_path.mkdir(parents=True, exist_ok=True)
+ store1.export_sidecar(index_path)
+
+ # Load into a fresh store
+ proc = MagicMock()
+ scores = [0.5, 0.8, 0.3]
+ proc.score.return_value = torch.tensor([scores])
+ store2 = LocalVectorStore()
+ store2.open("idx", tmp_path, create=False)
+ store2.load_sidecar(index_path)
+ store2.set_processor(proc)
+
+ # All points should be present
+ for p in pts:
+ assert store2.point_exists(p.point_id)
+
+ # Search should work
+ q = MultiVectorQuery(vectors=_dummy_embedding(2, 8))
+ results = store2.search(q, k=2)
+ assert len(results) == 2
diff --git a/tests/test_vector_store_milvus.py b/tests/test_vector_store_milvus.py
new file mode 100644
index 0000000..01699d5
--- /dev/null
+++ b/tests/test_vector_store_milvus.py
@@ -0,0 +1,339 @@
+"""Tests for MilvusVectorStore.
+
+Unit tests mock the MilvusClient.
+The slow integration test uses Milvus Lite (file-based, no external server needed).
+"""
+from __future__ import annotations
+
+import tempfile
+import uuid
+from pathlib import Path
+from unittest.mock import MagicMock, call, patch
+
+import pytest
+import torch
+
+from foretrieval.vector_store.base import (
+ MultiVectorQuery,
+ StoredPoint,
+ make_point_id,
+)
+from foretrieval.vector_store.milvus import (
+ MilvusVectorStore,
+ _MILVUS_AVAILABLE,
+ _collection_names,
+ _mean_pool,
+ _page_id_str,
+ _deserialize_payload,
+ _serialize_payload,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helper builders
+# ---------------------------------------------------------------------------
+
+def _point(doc_id: int, page_id: int, n_tokens: int = 4, dim: int = 8) -> StoredPoint:
+ pid = make_point_id(doc_id, page_id)
+ return StoredPoint(
+ point_id=pid,
+ vector=torch.rand(n_tokens, dim),
+ payload={"doc_id": doc_id, "page_id": page_id, "chunk_id": None, "metadata": {}},
+ )
+
+
+def _make_mock_store(index_name: str = "test_idx") -> tuple[MilvusVectorStore, MagicMock]:
+ store = MilvusVectorStore.__new__(MilvusVectorStore)
+ mock_client = MagicMock()
+ mock_client.list_collections.return_value = []
+ store._client = mock_client
+ store._index_name = index_name
+ store._db_path = Path("/tmp/fake.db")
+ store._candidate_limit = 64
+ return store, mock_client
+
+
+# ---------------------------------------------------------------------------
+# Pure helpers
+# ---------------------------------------------------------------------------
+
+class TestHelpers:
+ def test_mean_pool_correct_shape(self):
+ t = torch.ones(4, 8)
+ pool = _mean_pool(t)
+ assert len(pool) == 8
+
+ def test_mean_pool_values(self):
+ t = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+ pool = _mean_pool(t)
+ assert pool == pytest.approx([2.0, 3.0])
+
+ def test_serialize_deserialize_roundtrip(self):
+ payload = {"doc_id": 1, "page_id": 2, "metadata": {"cat": "A"}}
+ raw = _serialize_payload(payload)
+ result = _deserialize_payload(raw)
+ assert result["doc_id"] == 1
+
+ def test_collection_names_suffix(self):
+ page_col, token_col = _collection_names("myindex")
+ assert page_col == "myindex__pages"
+ assert token_col == "myindex__tokens"
+
+ def test_page_id_str(self):
+ pid = make_point_id(5, 3)
+ assert _page_id_str(pid) == str(pid)
+
+
+# ---------------------------------------------------------------------------
+# Optional-dependency guard
+# ---------------------------------------------------------------------------
+
+class TestOptionalDepGuard:
+ def test_missing_milvus_raises_on_open(self, tmp_path):
+ with patch("foretrieval.vector_store.milvus._MILVUS_AVAILABLE", False):
+ store = MilvusVectorStore()
+ with pytest.raises(RuntimeError, match="pymilvus"):
+ store.open("idx", tmp_path, create=True, dim=8)
+
+
+# ---------------------------------------------------------------------------
+# Collection management
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _MILVUS_AVAILABLE, reason="pymilvus not installed")
+class TestCollectionManagement:
+ def test_collection_not_exists_empty_list(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = []
+ assert not store.collection_exists()
+
+ def test_collection_exists_when_page_col_present(self):
+ store, mock_client = _make_mock_store()
+ page_col, _ = _collection_names("test_idx")
+ mock_client.list_collections.return_value = [page_col]
+ assert store.collection_exists()
+
+ def test_create_collection_creates_both_collections(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = []
+ mock_client.create_collection = MagicMock()
+ mock_client.prepare_index_params.return_value = MagicMock()
+ store.create_collection(dim=8)
+ assert mock_client.create_collection.call_count == 2
+
+
+# ---------------------------------------------------------------------------
+# Upsert
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _MILVUS_AVAILABLE, reason="pymilvus not installed")
+class TestUpsert:
+ def test_upsert_inserts_page_rows(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = [
+ "test_idx__pages", "test_idx__tokens"
+ ]
+ sp = _point(0, 1, n_tokens=3, dim=8)
+ store.upsert([sp])
+ page_col, token_col = _collection_names("test_idx")
+ upsert_calls = {c.kwargs["collection_name"]: c for c in mock_client.upsert.call_args_list}
+ assert page_col in upsert_calls
+
+ def test_upsert_inserts_token_rows(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = [
+ "test_idx__pages", "test_idx__tokens"
+ ]
+ n_tokens = 5
+ sp = _point(0, 1, n_tokens=n_tokens, dim=8)
+ store.upsert([sp])
+ page_col, token_col = _collection_names("test_idx")
+ upsert_calls = {c.kwargs["collection_name"]: c for c in mock_client.upsert.call_args_list}
+ assert token_col in upsert_calls
+ token_rows = upsert_calls[token_col].kwargs["data"]
+ assert len(token_rows) == n_tokens
+
+ def test_upsert_page_vector_is_mean_pooled(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = [
+ "test_idx__pages", "test_idx__tokens"
+ ]
+ vec = torch.ones(3, 4) # all ones → mean pool = [1,1,1,1]
+ pid = make_point_id(0, 1)
+ sp = StoredPoint(
+ point_id=pid,
+ vector=vec,
+ payload={"doc_id": 0, "page_id": 1, "chunk_id": None, "metadata": {}},
+ )
+ store.upsert([sp])
+ page_col, _ = _collection_names("test_idx")
+ upsert_calls = {c.kwargs["collection_name"]: c for c in mock_client.upsert.call_args_list}
+ page_row = upsert_calls[page_col].kwargs["data"][0]
+ assert page_row["page_vector"] == pytest.approx([1.0, 1.0, 1.0, 1.0])
+
+ def test_upsert_token_rows_share_page_id(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = [
+ "test_idx__pages", "test_idx__tokens"
+ ]
+ pid = make_point_id(0, 1)
+ sp = _point(0, 1, n_tokens=3)
+ store.upsert([sp])
+ _, token_col = _collection_names("test_idx")
+ upsert_calls = {c.kwargs["collection_name"]: c for c in mock_client.upsert.call_args_list}
+ token_rows = upsert_calls[token_col].kwargs["data"]
+ page_ids = {row["page_id"] for row in token_rows}
+ assert page_ids == {_page_id_str(pid)}
+
+ def test_point_exists_true(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = ["test_idx__pages"]
+ mock_client.get.return_value = [{"id": "100"}]
+ assert store.point_exists(100)
+
+ def test_point_exists_false(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = ["test_idx__pages"]
+ mock_client.get.return_value = []
+ assert not store.point_exists(100)
+
+
+# ---------------------------------------------------------------------------
+# Search — two-stage retrieval logic
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _MILVUS_AVAILABLE, reason="pymilvus not installed")
+class TestSearch:
+ def _make_store_with_candidates(self, candidate_ids: list[str]) -> tuple[MilvusVectorStore, MagicMock]:
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = [
+ "test_idx__pages", "test_idx__tokens"
+ ]
+
+ # _fetch_candidates: mock page search
+ def page_search(**kwargs):
+ rows = [
+ {
+ "id": cid,
+ "entity": {
+ "payload_json": _serialize_payload(
+ {"doc_id": int(cid) // 10_000_000, "page_id": 1, "metadata": {}}
+ )
+ },
+ }
+ for cid in candidate_ids
+ ]
+ return [rows]
+
+ # _late_interaction_rerank: mock token search, return first candidate as best
+ def token_search(**kwargs):
+ if not candidate_ids:
+ return [[]]
+ rows = [
+ {
+ "entity": {"page_id": candidate_ids[0]},
+ "distance": 0.9,
+ }
+ ]
+ return [rows]
+
+ mock_client.search.side_effect = [page_search(**{}), token_search(**{})]
+ return store, mock_client
+
+ def test_search_returns_search_hits(self):
+ candidate_ids = [str(make_point_id(i, 1)) for i in range(3)]
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = ["test_idx__pages", "test_idx__tokens"]
+
+ # Mock page search
+ page_rows = [
+ {"id": cid, "entity": {"payload_json": _serialize_payload({"doc_id": i, "page_id": 1, "metadata": {}})}}
+ for i, cid in enumerate(candidate_ids)
+ ]
+ # Mock token search (one per query token)
+ token_rows = [
+ {"entity": {"page_id": candidate_ids[0]}, "distance": 0.9}
+ ]
+ mock_client.search.side_effect = [[page_rows], [token_rows]]
+
+ q = MultiVectorQuery(vectors=torch.rand(1, 8))
+ results = store.search(q, k=2)
+ assert len(results) > 0
+ assert all(hasattr(r, "point_id") for r in results)
+
+ def test_search_empty_candidates_returns_empty(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = ["test_idx__pages", "test_idx__tokens"]
+ mock_client.search.return_value = [[]]
+ q = MultiVectorQuery(vectors=torch.rand(1, 8))
+ results = store.search(q, k=5)
+ assert results == []
+
+
+# ---------------------------------------------------------------------------
+# fetch_vector
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _MILVUS_AVAILABLE, reason="pymilvus not installed")
+class TestFetchVector:
+ def test_fetch_returns_tensor(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = ["test_idx__tokens"]
+ dim = 8
+ n_tokens = 3
+ token_rows = [{"token_vector": torch.rand(dim).tolist()} for _ in range(n_tokens)]
+ mock_client.query.return_value = token_rows
+ vec = store.fetch_vector(make_point_id(0, 1))
+ assert vec is not None
+ assert vec.shape == (n_tokens, dim)
+
+ def test_fetch_returns_none_when_no_rows(self):
+ store, mock_client = _make_mock_store()
+ mock_client.list_collections.return_value = ["test_idx__tokens"]
+ mock_client.query.return_value = []
+ assert store.fetch_vector(make_point_id(0, 1)) is None
+
+
+# ---------------------------------------------------------------------------
+# Integration test: Milvus Lite (no external server)
+# ---------------------------------------------------------------------------
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _MILVUS_AVAILABLE, reason="pymilvus not installed")
+class TestMilvusIntegration:
+ """Full round-trip using Milvus Lite (file-based)."""
+
+ def test_round_trip(self, tmp_path):
+ dim = 8
+ store = MilvusVectorStore()
+
+ # open with create=True but no dim → collection not created yet
+ store.open("test_rt", tmp_path, create=False)
+ assert not store.collection_exists()
+
+ # Create the collection explicitly
+ store.create_collection(dim)
+ assert store.collection_exists()
+
+ pts = [
+ StoredPoint(
+ point_id=make_point_id(i, 1),
+ vector=torch.rand(4, dim),
+ payload={"doc_id": i, "page_id": 1, "chunk_id": None, "metadata": {}},
+ )
+ for i in range(5)
+ ]
+ store.upsert(pts)
+
+ assert store.point_exists(pts[0].point_id)
+ assert not store.point_exists(make_point_id(99, 99))
+
+ q = MultiVectorQuery(vectors=torch.rand(2, dim))
+ results = store.search(q, k=3)
+ assert len(results) > 0
+
+ vec = store.fetch_vector(pts[0].point_id)
+ assert vec is not None
+ assert vec.shape[1] == dim
+
+ store.close()
diff --git a/tests/test_vector_store_qdrant.py b/tests/test_vector_store_qdrant.py
new file mode 100644
index 0000000..3dd447e
--- /dev/null
+++ b/tests/test_vector_store_qdrant.py
@@ -0,0 +1,268 @@
+"""Tests for QdrantVectorStore.
+
+Unit tests mock the QdrantClient.
+The slow integration test runs a real embedded Qdrant (no external server needed).
+"""
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch, PropertyMock
+
+import pytest
+import torch
+
+from foretrieval.vector_store.base import (
+ MultiVectorQuery,
+ StoredPoint,
+ make_point_id,
+)
+from foretrieval.vector_store.qdrant import QdrantVectorStore, _QDRANT_AVAILABLE
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _point(doc_id: int, page_id: int, dim: int = 8) -> StoredPoint:
+ pid = make_point_id(doc_id, page_id)
+ return StoredPoint(
+ point_id=pid,
+ vector=torch.rand(4, dim),
+ payload={"doc_id": doc_id, "page_id": page_id, "chunk_id": None, "metadata": {}},
+ )
+
+
+def _make_mock_store() -> tuple[QdrantVectorStore, MagicMock]:
+ """Return a QdrantVectorStore with a mock client already injected."""
+ store = QdrantVectorStore.__new__(QdrantVectorStore)
+ mock_client = MagicMock()
+ mock_client.collection_exists.return_value = False
+ store._client = mock_client
+ store._collection_name = "test_idx"
+ store._index_root = Path("/tmp/foretrieval_test")
+ return store, mock_client
+
+
+# ---------------------------------------------------------------------------
+# Optional-dependency guard
+# ---------------------------------------------------------------------------
+
+class TestOptionalDepGuard:
+ def test_missing_qdrant_raises_on_open(self, tmp_path):
+ with patch("foretrieval.vector_store.qdrant._QDRANT_AVAILABLE", False):
+ store = QdrantVectorStore()
+ with pytest.raises(RuntimeError, match="qdrant-client"):
+ store.open("idx", tmp_path, create=True, dim=8)
+
+
+# ---------------------------------------------------------------------------
+# collection_exists / create_collection
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _QDRANT_AVAILABLE, reason="qdrant-client not installed")
+class TestCollectionManagement:
+ def test_collection_exists_false_when_not_created(self):
+ store, mock_client = _make_mock_store()
+ mock_client.collection_exists.return_value = False
+ assert not store.collection_exists()
+
+ def test_collection_exists_true_when_created(self):
+ store, mock_client = _make_mock_store()
+ mock_client.collection_exists.return_value = True
+ assert store.collection_exists()
+
+ def test_create_collection_calls_client(self):
+ store, mock_client = _make_mock_store()
+ mock_client.collection_exists.return_value = False
+ store.create_collection(dim=128)
+ mock_client.create_collection.assert_called_once()
+ call_kwargs = mock_client.create_collection.call_args[1]
+ assert call_kwargs["collection_name"] == "test_idx"
+
+ def test_create_collection_multivector_config(self):
+ """Collection must be created with MultiVectorConfig(MAX_SIM)."""
+ from qdrant_client.models import MultiVectorComparator
+ store, mock_client = _make_mock_store()
+ mock_client.collection_exists.return_value = False
+ store.create_collection(dim=128)
+ call_kwargs = mock_client.create_collection.call_args[1]
+ vc = call_kwargs["vectors_config"]
+ assert vc.multivector_config.comparator == MultiVectorComparator.MAX_SIM
+
+ def test_create_collection_noop_if_already_exists(self):
+ store, mock_client = _make_mock_store()
+ mock_client.collection_exists.return_value = True
+ store.create_collection(dim=128)
+ mock_client.create_collection.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# Upsert
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _QDRANT_AVAILABLE, reason="qdrant-client not installed")
+class TestUpsert:
+ def test_upsert_calls_client_upsert(self):
+ store, mock_client = _make_mock_store()
+ sp = _point(0, 1)
+ store.upsert([sp])
+ mock_client.upsert.assert_called_once()
+ call_kwargs = mock_client.upsert.call_args[1]
+ assert call_kwargs["collection_name"] == "test_idx"
+ assert len(call_kwargs["points"]) == 1
+
+ def test_upsert_payload_contains_doc_page(self):
+ store, mock_client = _make_mock_store()
+ sp = _point(5, 3)
+ store.upsert([sp])
+ point = mock_client.upsert.call_args[1]["points"][0]
+ assert point.payload["doc_id"] == 5
+ assert point.payload["page_id"] == 3
+
+ def test_upsert_vector_is_list_of_floats(self):
+ store, mock_client = _make_mock_store()
+ sp = _point(0, 1)
+ store.upsert([sp])
+ point = mock_client.upsert.call_args[1]["points"][0]
+ # Multivector: list of lists
+ assert isinstance(point.vector, list)
+ assert all(isinstance(row, list) for row in point.vector)
+
+ def test_point_exists_returns_true(self):
+ store, mock_client = _make_mock_store()
+ mock_client.collection_exists.return_value = True
+ mock_client.retrieve.return_value = [MagicMock()]
+ assert store.point_exists(make_point_id(0, 1))
+
+ def test_point_exists_returns_false(self):
+ store, mock_client = _make_mock_store()
+ mock_client.collection_exists.return_value = True
+ mock_client.retrieve.return_value = []
+ assert not store.point_exists(make_point_id(0, 1))
+
+
+# ---------------------------------------------------------------------------
+# Search
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _QDRANT_AVAILABLE, reason="qdrant-client not installed")
+class TestSearch:
+ def _make_hit(self, pid, score, doc_id, page_id):
+ hit = MagicMock()
+ hit.id = pid
+ hit.score = score
+ hit.payload = {"doc_id": doc_id, "page_id": page_id, "chunk_id": None, "metadata": {}}
+ return hit
+
+ def test_search_returns_search_hits(self):
+ store, mock_client = _make_mock_store()
+ hits = [self._make_hit(100, 0.9, 0, 1), self._make_hit(200, 0.5, 1, 1)]
+ response = MagicMock()
+ response.points = hits
+ mock_client.query_points.return_value = response
+
+ q = MultiVectorQuery(vectors=torch.rand(2, 8))
+ results = store.search(q, k=2)
+ assert len(results) == 2
+ assert results[0].point_id == 100
+ assert results[0].score == pytest.approx(0.9)
+
+ def test_search_passes_k_to_client(self):
+ store, mock_client = _make_mock_store()
+ response = MagicMock()
+ response.points = []
+ mock_client.query_points.return_value = response
+ q = MultiVectorQuery(vectors=torch.rand(2, 8))
+ store.search(q, k=7)
+ call_kwargs = mock_client.query_points.call_args[1]
+ assert call_kwargs["limit"] == 7
+
+ def test_search_with_metadata_filter_passes_filter(self):
+ from qdrant_client.models import Filter
+ store, mock_client = _make_mock_store()
+ response = MagicMock()
+ response.points = []
+ mock_client.query_points.return_value = response
+ q = MultiVectorQuery(vectors=torch.rand(2, 8), filter_metadata={"category": "A"})
+ store.search(q, k=5)
+ call_kwargs = mock_client.query_points.call_args[1]
+ assert call_kwargs["query_filter"] is not None
+ assert isinstance(call_kwargs["query_filter"], Filter)
+
+ def test_search_no_filter_passes_none(self):
+ store, mock_client = _make_mock_store()
+ response = MagicMock()
+ response.points = []
+ mock_client.query_points.return_value = response
+ q = MultiVectorQuery(vectors=torch.rand(2, 8))
+ store.search(q, k=5)
+ call_kwargs = mock_client.query_points.call_args[1]
+ assert call_kwargs["query_filter"] is None
+
+
+# ---------------------------------------------------------------------------
+# fetch_vector
+# ---------------------------------------------------------------------------
+
+@pytest.mark.skipif(not _QDRANT_AVAILABLE, reason="qdrant-client not installed")
+class TestFetchVector:
+ def test_fetch_returns_tensor(self):
+ store, mock_client = _make_mock_store()
+ vec_data = [[0.1, 0.2, 0.3, 0.4]] * 4
+ retrieved = MagicMock()
+ retrieved.vector = vec_data
+ mock_client.retrieve.return_value = [retrieved]
+ result = store.fetch_vector(make_point_id(0, 1))
+ assert result is not None
+ assert isinstance(result, torch.Tensor)
+
+ def test_fetch_returns_none_when_not_found(self):
+ store, mock_client = _make_mock_store()
+ mock_client.retrieve.return_value = []
+ assert store.fetch_vector(make_point_id(99, 99)) is None
+
+
+# ---------------------------------------------------------------------------
+# Integration test: real embedded Qdrant
+# ---------------------------------------------------------------------------
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _QDRANT_AVAILABLE, reason="qdrant-client not installed")
+class TestQdrantIntegration:
+ """Full round-trip: open → upsert → search → fetch_vector → close."""
+
+ def test_round_trip(self, tmp_path):
+ dim = 8
+ store = QdrantVectorStore()
+
+ # Open with create=False first to test empty state
+ store.open("test_rt", tmp_path, create=False)
+ assert not store.collection_exists()
+
+ # Create collection explicitly
+ store.create_collection(dim)
+ assert store.collection_exists()
+
+ pts = [
+ StoredPoint(
+ point_id=make_point_id(i, 1),
+ vector=torch.rand(4, dim),
+ payload={"doc_id": i, "page_id": 1, "chunk_id": None, "metadata": {}},
+ )
+ for i in range(5)
+ ]
+ store.upsert(pts)
+
+ assert store.point_exists(pts[0].point_id)
+ assert not store.point_exists(make_point_id(99, 99))
+
+ q = MultiVectorQuery(vectors=torch.rand(2, dim))
+ results = store.search(q, k=3)
+ assert len(results) == 3
+
+ vec = store.fetch_vector(pts[0].point_id)
+ assert vec is not None
+ assert vec.shape[1] == dim
+
+ store.close()
diff --git a/tests/test_vector_store_remote.py b/tests/test_vector_store_remote.py
new file mode 100644
index 0000000..5d5addb
--- /dev/null
+++ b/tests/test_vector_store_remote.py
@@ -0,0 +1,208 @@
+"""Tests for RemoteVectorStore — all HTTP calls mocked via a MagicMock client."""
+
+from pathlib import Path
+from unittest.mock import MagicMock
+import pytest
+import torch
+
+from foretrieval.vector_store.remote import RemoteVectorStore
+from foretrieval.vector_store.base import MultiVectorQuery, SearchHit, StoredPoint, make_point_id
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_mock_client() -> MagicMock:
+ client = MagicMock()
+ client.open_collection.return_value = {"opened": True, "backend": "qdrant", "created": True}
+ client.collection_exists.return_value = True
+ client.point_exists.return_value = False
+ client.search.return_value = []
+ client.fetch_vector.return_value = None
+ return client
+
+
+def _make_store(backend="qdrant") -> tuple[RemoteVectorStore, MagicMock]:
+ client = _make_mock_client()
+ store = RemoteVectorStore(client, backend=backend)
+ return store, client
+
+
+def _make_point(doc_id=1, page_id=0, dim=8) -> StoredPoint:
+ return StoredPoint(
+ point_id=make_point_id(doc_id, page_id),
+ vector=torch.randn(3, dim),
+ payload={"doc_id": doc_id, "page_id": page_id},
+ )
+
+
+# ---------------------------------------------------------------------------
+# Lifecycle
+# ---------------------------------------------------------------------------
+
+class TestLifecycle:
+ def test_open_calls_client(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True, dim=128)
+ client.open_collection.assert_called_once_with(
+ "idx", "qdrant", create=True, dim=128, storage_config=None
+ )
+
+ def test_is_opened_after_open(self):
+ store, _ = _make_store()
+ assert not store.is_opened
+ store.open("idx", Path("."), create=True, dim=128)
+ assert store.is_opened
+
+ def test_close_marks_not_opened(self):
+ store, _ = _make_store()
+ store.open("idx", Path("."), create=True, dim=128)
+ store.close()
+ assert not store.is_opened
+
+ def test_backend_name_class_attr(self):
+ assert RemoteVectorStore.backend_name == "remote"
+
+
+# ---------------------------------------------------------------------------
+# collection_exists
+# ---------------------------------------------------------------------------
+
+class TestCollectionExists:
+ def test_returns_false_before_open(self):
+ store, client = _make_store()
+ assert store.collection_exists() is False
+
+ def test_delegates_to_client(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ client.collection_exists.return_value = True
+ assert store.collection_exists() is True
+ client.collection_exists.assert_called_with("idx")
+
+
+# ---------------------------------------------------------------------------
+# upsert / point_exists
+# ---------------------------------------------------------------------------
+
+class TestWrite:
+ def test_upsert_delegates_to_client(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ pts = [_make_point()]
+ store.upsert(pts)
+ client.upsert.assert_called_once_with("idx", pts)
+
+ def test_point_exists_delegates(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ client.point_exists.return_value = True
+ pid = make_point_id(1, 0)
+ assert store.point_exists(pid) is True
+ client.point_exists.assert_called_with("idx", pid)
+
+ def test_point_exists_returns_false_before_open(self):
+ store, _ = _make_store()
+ assert store.point_exists(12345) is False
+
+ def test_upsert_raises_without_open(self):
+ store, _ = _make_store()
+ with pytest.raises(RuntimeError, match="open"):
+ store.upsert([_make_point()])
+
+
+# ---------------------------------------------------------------------------
+# search
+# ---------------------------------------------------------------------------
+
+class TestSearch:
+ def test_search_delegates_to_client(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ hit = SearchHit(point_id=10000, score=9.5, payload={"doc_id": 1, "page_id": 0})
+ client.search.return_value = [hit]
+
+ query = MultiVectorQuery(vectors=torch.randn(4, 8))
+ results = store.search(query, k=3)
+
+ client.search.assert_called_once_with("idx", query, 3)
+ assert len(results) == 1
+ assert results[0].score == pytest.approx(9.5)
+
+ def test_search_raises_without_open(self):
+ store, _ = _make_store()
+ with pytest.raises(RuntimeError, match="open"):
+ store.search(MultiVectorQuery(vectors=torch.randn(2, 8)), k=5)
+
+
+# ---------------------------------------------------------------------------
+# fetch_vector
+# ---------------------------------------------------------------------------
+
+class TestFetchVector:
+ def test_fetch_returns_tensor(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ t = torch.randn(3, 8)
+ client.fetch_vector.return_value = t
+ result = store.fetch_vector(make_point_id(1, 0))
+ assert result is t
+
+ def test_fetch_returns_none_on_miss(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ client.fetch_vector.return_value = None
+ assert store.fetch_vector(99999) is None
+
+ def test_fetch_returns_none_before_open(self):
+ store, _ = _make_store()
+ assert store.fetch_vector(12345) is None
+
+
+# ---------------------------------------------------------------------------
+# Persistence no-ops
+# ---------------------------------------------------------------------------
+
+class TestPersistenceNoOps:
+ def test_export_sidecar_noop(self, tmp_path):
+ store, _ = _make_store()
+ store.open("idx", Path("."), create=True)
+ # Should not raise or create any files
+ store.export_sidecar(tmp_path)
+ assert list(tmp_path.iterdir()) == []
+
+ def test_load_sidecar_noop(self, tmp_path):
+ store, _ = _make_store()
+ store.open("idx", Path("."), create=True)
+ store.load_sidecar(tmp_path) # no error
+
+
+# ---------------------------------------------------------------------------
+# Bookkeeping (server-side, replaces local sidecars in remote mode)
+# ---------------------------------------------------------------------------
+
+class TestBookkeeping:
+ def test_supports_remote_bookkeeping(self):
+ store, _ = _make_store()
+ assert store.supports_remote_bookkeeping() is True
+
+ def test_export_bookkeeping_delegates(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ blob = {"index_config": {"model_name": "m"}}
+ store.export_bookkeeping(blob)
+ client.put_bookkeeping.assert_called_once_with("idx", blob)
+
+ def test_load_bookkeeping_delegates(self):
+ store, client = _make_store()
+ store.open("idx", Path("."), create=True)
+ client.get_bookkeeping.return_value = {"index_config": {"model_name": "m"}}
+ out = store.load_bookkeeping()
+ client.get_bookkeeping.assert_called_once_with("idx")
+ assert out["index_config"]["model_name"] == "m"
+
+ def test_export_bookkeeping_before_open_raises(self):
+ store, _ = _make_store()
+ with pytest.raises(RuntimeError):
+ store.export_bookkeeping({})
![\"A](\"./docs/product_c.png\")
"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 14,
- "id": "e8a931d6-f28c-496c-a0c4-31bd4eb0ffd5",
- "metadata": {},
- "outputs": [],
- "source": [
- "query = \"In which month did Product C generate the most revenue?\""
- ]
- },
- {
- "cell_type": "markdown",
- "id": "7fd8752a-951a-4763-baff-9a09bfc1b928",
- "metadata": {},
- "source": [
- "We're going to go through the same process as for the Attention is All You Need paper, and we'll index our financial report:"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 15,
- "id": "65c2a13c-ee4e-4b45-8eb9-f466b5b2b93d",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Added page 1 of document 1 to index.\n",
- "Added page 2 of document 1 to index.\n",
- "Added page 3 of document 1 to index.\n",
- "Added page 4 of document 1 to index.\n",
- "Added page 5 of document 1 to index.\n",
- "Added page 6 of document 1 to index.\n",
- "Index exported to .byaldi/financial_report\n",
- "Index exported to .byaldi/financial_report\n"
- ]
- },
- {
- "data": {
- "text/plain": [
- "{0: 'docs/attention.pdf', 1: 'docs/financial_report.pdf'}"
- ]
- },
- "execution_count": 15,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "RAG.index(\n",
- " input_path=\"./docs/financial_report.pdf\",\n",
- " index_name=\"financial_report\",\n",
- " store_collection_with_index=True,\n",
- " overwrite=True\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "55cc7aa6-baf3-4297-ad47-3e404702009f",
- "metadata": {},
- "source": [
- "Now, let's search the index for our query. We expect the top result to be page 4."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 16,
- "id": "a3d9fec7-b262-41a4-a3e4-82b7c7356983",
- "metadata": {},
- "outputs": [
- {
- "data": {
- "text/plain": [
- "4"
- ]
- },
- "execution_count": 16,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "results = RAG.search(query, k=1)\n",
- "results[0].page_num"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "2218061e-57fd-477a-8812-d206db602e6c",
- "metadata": {},
- "source": [
- "Finally, we again repeat the process we went through earlier: first, convert the top search result to bytes, then pass it to Claude with our query. If everything works as expected Claude should tell us Product C generated the most revenue in **June**."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 17,
- "id": "8598f8af-a701-45c5-9c48-f214ef5bd4ff",
- "metadata": {},
- "outputs": [
- {
- "data": {
- "text/markdown": [
- "According to the bar graph showing monthly revenue for Product C, the month with the highest revenue was June. The bar for June is visibly the tallest, reaching above 2500 on the revenue scale, indicating it generated the most revenue compared to all other months shown.\n",
- "\n",
- "