diff --git a/CHANGELOG.md b/CHANGELOG.md index 84674d9..94586c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this - Soft-delete chunks: delete button with confirmation dialog, chunks hidden from UI but preserved in data - Vector index metadata schema: `IndexedChunk` domain model, OpenSearch mapping builder, configurable embedding dimension - `VectorStore` port (Protocol): `ensure_index`, `index_chunks`, `search_similar`, `get_chunks`, `delete_document` +- OpenSearch adapter (`OpenSearchStore`): kNN vector search, full-text search, bulk indexing, document CRUD ### Fixed diff --git a/document-parser/domain/ports.py b/document-parser/domain/ports.py index aa852cc..0375da5 100644 --- a/document-parser/domain/ports.py +++ b/document-parser/domain/ports.py @@ -6,7 +6,7 @@ Infrastructure adapters (local Docling, Docling Serve, etc.) implement these. from __future__ import annotations -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, Protocol, runtime_checkable if TYPE_CHECKING: from domain.models import AnalysisJob, Document @@ -82,6 +82,7 @@ class AnalysisRepository(Protocol): async def delete_by_document(self, document_id: str) -> int: ... +@runtime_checkable class VectorStore(Protocol): """Port for vector storage and retrieval. diff --git a/document-parser/infra/opensearch_store.py b/document-parser/infra/opensearch_store.py new file mode 100644 index 0000000..c4a97f9 --- /dev/null +++ b/document-parser/infra/opensearch_store.py @@ -0,0 +1,204 @@ +"""OpenSearch adapter implementing the VectorStore port. + +Uses the opensearch-py client for kNN vector search, full-text search, +and document CRUD against an OpenSearch cluster. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from opensearchpy import AsyncOpenSearch, NotFoundError + +from domain.vector_schema import ( + ChunkBboxEntry, + ChunkOrigin, + DocItemRef, + IndexedChunk, + SearchResult, +) + +logger = logging.getLogger(__name__) + + +def _hit_to_indexed_chunk(hit: dict[str, Any]) -> IndexedChunk: + """Reconstruct an IndexedChunk from an OpenSearch _source document.""" + src = hit["_source"] + origin_raw = src.get("origin") + origin = ( + ChunkOrigin(binary_hash=origin_raw["binary_hash"], filename=origin_raw["filename"]) + if origin_raw + else None + ) + return IndexedChunk( + doc_id=src["doc_id"], + filename=src["filename"], + content=src["content"], + embedding=src.get("embedding", []), + chunk_index=src["chunk_index"], + chunk_type=src["chunk_type"], + page_number=src["page_number"], + bboxes=[ + ChunkBboxEntry(page=b["page"], x=b["x"], y=b["y"], w=b["w"], h=b["h"]) + for b in src.get("bboxes", []) + ], + headings=src.get("headings", []), + doc_items=[ + DocItemRef(self_ref=d["self_ref"], label=d["label"]) for d in src.get("doc_items", []) + ], + origin=origin, + ) + + +def _hit_to_result(hit: dict[str, Any]) -> SearchResult: + """Convert an OpenSearch hit to a SearchResult.""" + return SearchResult( + chunk=_hit_to_indexed_chunk(hit), + score=hit.get("_score", 0.0), + ) + + +class OpenSearchStore: + """Concrete VectorStore adapter backed by OpenSearch. + + Satisfies the ``VectorStore`` Protocol defined in ``domain.ports``. + + Args: + url: OpenSearch cluster URL (e.g. ``http://localhost:9200``). + verify_certs: Whether to verify TLS certificates. + """ + + def __init__(self, url: str, *, verify_certs: bool = False) -> None: + self._client = AsyncOpenSearch( + hosts=[url], + use_ssl=url.startswith("https"), + verify_certs=verify_certs, + ssl_show_warn=False, + ) + + # -- lifecycle ------------------------------------------------------------- + + async def close(self) -> None: + """Close the underlying HTTP connection pool.""" + await self._client.close() + + # -- VectorStore protocol methods ------------------------------------------ + + async def ensure_index(self, index_name: str, mapping: dict) -> None: + """Create the index if it does not exist. No-op if it already exists.""" + exists = await self._client.indices.exists(index=index_name) + if not exists: + await self._client.indices.create(index=index_name, body=mapping) + logger.info("Created OpenSearch index '%s'", index_name) + else: + logger.debug("Index '%s' already exists — skipping creation", index_name) + + async def index_chunks(self, index_name: str, chunks: list[IndexedChunk]) -> int: + """Bulk-index a list of chunks. Returns the number successfully indexed.""" + if not chunks: + return 0 + + body: list[dict[str, Any]] = [] + for chunk in chunks: + doc_id = f"{chunk.doc_id}_{chunk.chunk_index}" + body.append({"index": {"_index": index_name, "_id": doc_id}}) + body.append(chunk.to_dict()) + + resp = await self._client.bulk(body=body, refresh="wait_for") + + errors = sum(1 for item in resp["items"] if item["index"].get("error")) + indexed = len(chunks) - errors + if errors: + logger.warning("Bulk index to '%s': %d/%d failed", index_name, errors, len(chunks)) + return indexed + + async def search_similar( + self, + index_name: str, + embedding: list[float], + *, + k: int = 10, + doc_id: str | None = None, + ) -> list[SearchResult]: + """kNN search for the k nearest chunks by embedding similarity.""" + knn_query: dict[str, Any] = { + "knn": { + "embedding": { + "vector": embedding, + "k": k, + }, + }, + } + if doc_id: + knn_query["knn"]["embedding"]["filter"] = { + "term": {"doc_id": doc_id}, + } + + resp = await self._client.search( + index=index_name, + body={"size": k, "query": knn_query}, + _source_excludes=["embedding"], + ) + return [_hit_to_result(hit) for hit in resp["hits"]["hits"]] + + async def get_chunks( + self, + index_name: str, + doc_id: str, + *, + limit: int = 1000, + ) -> list[SearchResult]: + """Retrieve all indexed chunks for a document, ordered by chunk_index.""" + resp = await self._client.search( + index=index_name, + body={ + "size": limit, + "query": {"term": {"doc_id": doc_id}}, + "sort": [{"chunk_index": {"order": "asc"}}], + }, + _source_excludes=["embedding"], + ) + return [_hit_to_result(hit) for hit in resp["hits"]["hits"]] + + async def delete_document(self, index_name: str, doc_id: str) -> int: + """Delete all chunks for a document. Returns the number deleted.""" + try: + resp = await self._client.delete_by_query( + index=index_name, + body={"query": {"term": {"doc_id": doc_id}}}, + refresh=True, + ) + deleted: int = resp.get("deleted", 0) + return deleted + except NotFoundError: + return 0 + + # -- full-text search (bonus from spec) ------------------------------------ + + async def search_fulltext( + self, + index_name: str, + query_text: str, + *, + k: int = 10, + doc_id: str | None = None, + ) -> list[SearchResult]: + """Full-text search on the content field. + + This method is not part of the VectorStore protocol but is specified + in the issue acceptance criteria. + """ + must: list[dict[str, Any]] = [{"match": {"content": query_text}}] + if doc_id: + must.append({"term": {"doc_id": doc_id}}) + + resp = await self._client.search( + index=index_name, + body={ + "size": k, + "query": {"bool": {"must": must}}, + }, + _source_excludes=["embedding"], + ) + return [_hit_to_result(hit) for hit in resp["hits"]["hits"]] diff --git a/document-parser/requirements.txt b/document-parser/requirements.txt index 67cae25..81f72a7 100644 --- a/document-parser/requirements.txt +++ b/document-parser/requirements.txt @@ -7,3 +7,4 @@ pillow>=10.0.0,<11.0.0 aiosqlite>=0.20.0,<1.0.0 httpx>=0.27.0,<1.0.0 pypdfium2>=4.0.0,<5.0.0 +opensearch-py[async]>=2.6.0,<3.0.0 diff --git a/document-parser/tests/test_opensearch_store.py b/document-parser/tests/test_opensearch_store.py new file mode 100644 index 0000000..e29bf0a --- /dev/null +++ b/document-parser/tests/test_opensearch_store.py @@ -0,0 +1,320 @@ +"""Tests for the OpenSearch adapter (infra.opensearch_store). + +These tests mock the AsyncOpenSearch client to validate adapter logic +without requiring a running OpenSearch instance. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from domain.ports import VectorStore +from domain.vector_schema import ( + ChunkBboxEntry, + ChunkOrigin, + DocItemRef, + IndexedChunk, + SearchResult, + build_index_mapping, +) +from infra.opensearch_store import OpenSearchStore, _hit_to_indexed_chunk, _hit_to_result + +# -- Fixtures ----------------------------------------------------------------- + + +def _make_chunk( + doc_id: str = "doc-1", + chunk_index: int = 0, + content: str = "hello world", + embedding: list[float] | None = None, +) -> IndexedChunk: + return IndexedChunk( + doc_id=doc_id, + filename="test.pdf", + content=content, + embedding=embedding or [0.1, 0.2, 0.3], + chunk_index=chunk_index, + chunk_type="text", + page_number=1, + bboxes=[ChunkBboxEntry(page=1, x=0.0, y=0.0, w=100.0, h=50.0)], + headings=["Chapter 1"], + doc_items=[DocItemRef(self_ref="#/texts/0", label="text")], + origin=ChunkOrigin(binary_hash="abc123", filename="test.pdf"), + ) + + +def _make_hit( + doc_id: str = "doc-1", + chunk_index: int = 0, + score: float = 0.95, + content: str = "hello world", +) -> dict: + return { + "_id": f"{doc_id}_{chunk_index}", + "_score": score, + "_source": { + "doc_id": doc_id, + "filename": "test.pdf", + "content": content, + "chunk_index": chunk_index, + "chunk_type": "text", + "page_number": 1, + "bboxes": [{"page": 1, "x": 0.0, "y": 0.0, "w": 100.0, "h": 50.0}], + "headings": ["Chapter 1"], + "doc_items": [{"self_ref": "#/texts/0", "label": "text"}], + "origin": {"binary_hash": "abc123", "filename": "test.pdf"}, + }, + } + + +@pytest.fixture +def store() -> OpenSearchStore: + return OpenSearchStore("http://localhost:9200") + + +@pytest.fixture +def mock_client(store: OpenSearchStore) -> AsyncMock: + client = AsyncMock() + store._client = client + return client + + +# -- Protocol satisfaction ----------------------------------------------------- + + +class TestProtocolSatisfaction: + def test_satisfies_vector_store_protocol(self) -> None: + """OpenSearchStore structurally satisfies VectorStore Protocol.""" + store = OpenSearchStore("http://localhost:9200") + assert isinstance(store, VectorStore) + + +# -- Hit deserialization ------------------------------------------------------- + + +class TestHitDeserialization: + def test_hit_to_indexed_chunk(self) -> None: + hit = _make_hit() + chunk = _hit_to_indexed_chunk(hit) + assert isinstance(chunk, IndexedChunk) + assert chunk.doc_id == "doc-1" + assert chunk.content == "hello world" + assert chunk.chunk_index == 0 + assert chunk.page_number == 1 + assert len(chunk.bboxes) == 1 + assert chunk.bboxes[0].w == 100.0 + assert len(chunk.doc_items) == 1 + assert chunk.doc_items[0].label == "text" + assert chunk.origin is not None + assert chunk.origin.binary_hash == "abc123" + + def test_hit_to_indexed_chunk_no_origin(self) -> None: + hit = _make_hit() + hit["_source"]["origin"] = None + chunk = _hit_to_indexed_chunk(hit) + assert chunk.origin is None + + def test_hit_to_indexed_chunk_missing_optional_fields(self) -> None: + hit = _make_hit() + del hit["_source"]["bboxes"] + del hit["_source"]["headings"] + del hit["_source"]["doc_items"] + del hit["_source"]["origin"] + chunk = _hit_to_indexed_chunk(hit) + assert chunk.bboxes == [] + assert chunk.headings == [] + assert chunk.doc_items == [] + assert chunk.origin is None + + def test_hit_to_result(self) -> None: + hit = _make_hit(score=0.87) + result = _hit_to_result(hit) + assert isinstance(result, SearchResult) + assert result.score == 0.87 + assert result.chunk.doc_id == "doc-1" + + +# -- ensure_index -------------------------------------------------------------- + + +class TestEnsureIndex: + async def test_creates_index_when_not_exists( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + mock_client.indices.exists.return_value = False + mapping = build_index_mapping() + await store.ensure_index("test-index", mapping) + mock_client.indices.create.assert_awaited_once_with(index="test-index", body=mapping) + + async def test_noop_when_index_exists( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + mock_client.indices.exists.return_value = True + await store.ensure_index("test-index", {}) + mock_client.indices.create.assert_not_awaited() + + +# -- index_chunks -------------------------------------------------------------- + + +class TestIndexChunks: + async def test_bulk_indexes_chunks( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + chunks = [_make_chunk(chunk_index=0), _make_chunk(chunk_index=1)] + mock_client.bulk.return_value = { + "errors": False, + "items": [ + {"index": {"_id": "doc-1_0", "status": 201}}, + {"index": {"_id": "doc-1_1", "status": 201}}, + ], + } + count = await store.index_chunks("test-index", chunks) + assert count == 2 + mock_client.bulk.assert_awaited_once() + + async def test_returns_zero_for_empty_list( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + count = await store.index_chunks("test-index", []) + assert count == 0 + mock_client.bulk.assert_not_awaited() + + async def test_counts_partial_failures( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + chunks = [_make_chunk(chunk_index=0), _make_chunk(chunk_index=1)] + mock_client.bulk.return_value = { + "errors": True, + "items": [ + {"index": {"_id": "doc-1_0", "status": 201}}, + {"index": {"_id": "doc-1_1", "error": {"reason": "mapping"}}}, + ], + } + count = await store.index_chunks("test-index", chunks) + assert count == 1 + + async def test_bulk_body_structure( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + chunk = _make_chunk(doc_id="d1", chunk_index=3) + mock_client.bulk.return_value = { + "errors": False, + "items": [{"index": {"_id": "d1_3", "status": 201}}], + } + await store.index_chunks("idx", [chunk]) + call_body = mock_client.bulk.call_args[1]["body"] + assert call_body[0] == {"index": {"_index": "idx", "_id": "d1_3"}} + assert call_body[1]["doc_id"] == "d1" + assert call_body[1]["chunk_index"] == 3 + + +# -- search_similar ------------------------------------------------------------ + + +class TestSearchSimilar: + async def test_knn_search(self, store: OpenSearchStore, mock_client: AsyncMock) -> None: + mock_client.search.return_value = {"hits": {"hits": [_make_hit(score=0.99)]}} + results = await store.search_similar("idx", [0.1, 0.2, 0.3], k=5) + assert len(results) == 1 + assert results[0].score == 0.99 + call_body = mock_client.search.call_args[1]["body"] + assert call_body["query"]["knn"]["embedding"]["vector"] == [0.1, 0.2, 0.3] + assert call_body["query"]["knn"]["embedding"]["k"] == 5 + + async def test_knn_search_with_doc_filter( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + mock_client.search.return_value = {"hits": {"hits": []}} + await store.search_similar("idx", [0.1], doc_id="doc-42") + call_body = mock_client.search.call_args[1]["body"] + assert call_body["query"]["knn"]["embedding"]["filter"] == {"term": {"doc_id": "doc-42"}} + + async def test_knn_search_no_filter_by_default( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + mock_client.search.return_value = {"hits": {"hits": []}} + await store.search_similar("idx", [0.1]) + call_body = mock_client.search.call_args[1]["body"] + assert "filter" not in call_body["query"]["knn"]["embedding"] + + +# -- get_chunks ---------------------------------------------------------------- + + +class TestGetChunks: + async def test_retrieves_by_doc_id( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + mock_client.search.return_value = { + "hits": {"hits": [_make_hit(chunk_index=0), _make_hit(chunk_index=1)]} + } + results = await store.get_chunks("idx", "doc-1") + assert len(results) == 2 + call_body = mock_client.search.call_args[1]["body"] + assert call_body["query"] == {"term": {"doc_id": "doc-1"}} + assert call_body["sort"] == [{"chunk_index": {"order": "asc"}}] + + async def test_respects_limit(self, store: OpenSearchStore, mock_client: AsyncMock) -> None: + mock_client.search.return_value = {"hits": {"hits": []}} + await store.get_chunks("idx", "doc-1", limit=50) + call_body = mock_client.search.call_args[1]["body"] + assert call_body["size"] == 50 + + +# -- delete_document ----------------------------------------------------------- + + +class TestDeleteDocument: + async def test_deletes_by_doc_id(self, store: OpenSearchStore, mock_client: AsyncMock) -> None: + mock_client.delete_by_query.return_value = {"deleted": 5} + count = await store.delete_document("idx", "doc-1") + assert count == 5 + call_body = mock_client.delete_by_query.call_args[1]["body"] + assert call_body["query"] == {"term": {"doc_id": "doc-1"}} + + async def test_returns_zero_on_not_found( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + from opensearchpy import NotFoundError + + mock_client.delete_by_query.side_effect = NotFoundError(404, "index_not_found") + count = await store.delete_document("idx", "doc-1") + assert count == 0 + + +# -- search_fulltext ----------------------------------------------------------- + + +class TestSearchFulltext: + async def test_fulltext_search(self, store: OpenSearchStore, mock_client: AsyncMock) -> None: + mock_client.search.return_value = { + "hits": {"hits": [_make_hit(content="matching text", score=1.5)]} + } + results = await store.search_fulltext("idx", "matching") + assert len(results) == 1 + assert results[0].chunk.content == "matching text" + call_body = mock_client.search.call_args[1]["body"] + assert {"match": {"content": "matching"}} in call_body["query"]["bool"]["must"] + + async def test_fulltext_search_with_doc_filter( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + mock_client.search.return_value = {"hits": {"hits": []}} + await store.search_fulltext("idx", "query", doc_id="doc-5") + call_body = mock_client.search.call_args[1]["body"] + must_clauses = call_body["query"]["bool"]["must"] + assert {"term": {"doc_id": "doc-5"}} in must_clauses + + +# -- close --------------------------------------------------------------------- + + +class TestClose: + async def test_close_delegates_to_client( + self, store: OpenSearchStore, mock_client: AsyncMock + ) -> None: + await store.close() + mock_client.close.assert_awaited_once()