diff --git a/.env.example b/.env.example index b423dce..d7f515c 100644 --- a/.env.example +++ b/.env.example @@ -48,3 +48,8 @@ # Embedding vector dimension (default: 384 for Granite Embedding 30M / all-MiniLM-L6-v2) # EMBEDDING_DIMENSION=384 + +# Neo4j — graph-native document structure (used by docker-compose.dev.yml) +# NEO4J_URI=bolt://neo4j:7687 +# NEO4J_USER=neo4j +# NEO4J_PASSWORD=changeme diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 9255c3b..9a23fff 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -9,6 +9,26 @@ # ============================================================================= services: + # --- Neo4j (graph-native document structure) --- + neo4j: + image: neo4j:5.15-community + environment: + NEO4J_AUTH: ${NEO4J_USER:-neo4j}/${NEO4J_PASSWORD:-changeme} + NEO4J_server_memory_heap_initial__size: 512m + NEO4J_server_memory_heap_max__size: 1g + ports: + - "7474:7474" + - "7687:7687" + volumes: + - neo4j_data:/data + - neo4j_logs:/logs + healthcheck: + test: ["CMD-SHELL", "cypher-shell -u $${NEO4J_USER:-neo4j} -p $${NEO4J_PASSWORD:-changeme} 'RETURN 1' || exit 1"] + interval: 10s + timeout: 5s + retries: 10 + start_period: 30s + # --- OpenSearch (single-node, security disabled for local dev) --- opensearch: image: opensearchproject/opensearch:2 @@ -77,12 +97,17 @@ services: BATCH_PAGE_SIZE: ${BATCH_PAGE_SIZE:-10} OPENSEARCH_URL: http://opensearch:9200 EMBEDDING_URL: http://embedding:8001 + NEO4J_URI: bolt://neo4j:7687 + NEO4J_USER: ${NEO4J_USER:-neo4j} + NEO4J_PASSWORD: ${NEO4J_PASSWORD:-changeme} command: ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"] depends_on: opensearch: condition: service_healthy embedding: condition: service_healthy + neo4j: + condition: service_healthy deploy: resources: limits: @@ -105,6 +130,8 @@ services: volumes: opensearch_data: + neo4j_data: + neo4j_logs: uploads_data: db_data: frontend_node_modules: diff --git a/docker-compose.yml b/docker-compose.yml index 918bfb9..49c9880 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,4 +1,22 @@ services: + # --- Neo4j (graph-native document structure) --- + neo4j: + profiles: ["ingestion"] + image: neo4j:5.15-community + environment: + NEO4J_AUTH: ${NEO4J_USER:-neo4j}/${NEO4J_PASSWORD:-changeme} + NEO4J_server_memory_heap_initial__size: 512m + NEO4J_server_memory_heap_max__size: 1g + volumes: + - neo4j_data:/data + - neo4j_logs:/logs + healthcheck: + test: ["CMD-SHELL", "cypher-shell -u $${NEO4J_USER:-neo4j} -p $${NEO4J_PASSWORD:-changeme} 'RETURN 1' || exit 1"] + interval: 10s + timeout: 5s + retries: 10 + start_period: 30s + # --- OpenSearch (single-node, security disabled) --- opensearch: profiles: ["ingestion"] @@ -53,6 +71,9 @@ services: BATCH_PAGE_SIZE: ${BATCH_PAGE_SIZE:-10} OPENSEARCH_URL: ${OPENSEARCH_URL:-} EMBEDDING_URL: ${EMBEDDING_URL:-} + NEO4J_URI: ${NEO4J_URI:-} + NEO4J_USER: ${NEO4J_USER:-neo4j} + NEO4J_PASSWORD: ${NEO4J_PASSWORD:-changeme} deploy: resources: limits: @@ -69,5 +90,7 @@ services: volumes: opensearch_data: + neo4j_data: + neo4j_logs: uploads_data: db_data: diff --git a/docs/design/neo4j-integration.md b/docs/design/neo4j-integration.md new file mode 100644 index 0000000..1965408 --- /dev/null +++ b/docs/design/neo4j-integration.md @@ -0,0 +1,435 @@ +# Neo4j integration — Docling-Studio v0.5.0 + +Design doc for Neo4j integration targeting release 0.5.0. +Target: Hackernoon hackathon demo (Neo4j partner). + +--- + +## 1. Context and goals + +### Already in Docling-Studio +- Ingestion pipeline: Docling parser → chunking (HybridChunker) → embedding → OpenSearch (vector index) +- Vue 3 + FastAPI UI +- Debug view to inspect/edit chunks before retrieval +- Docker compose with existing services + +### What we add in v0.5.0 +- Neo4j as **graph-native storage** of the document structure +- A new ingestion layer that stores the DoclingDocument tree faithfully as a graph +- Minimal UI to visualize the graph (demo value to the judges) +- Compose pipeline with Neo4j + +### Why graph-native (hackathon positioning) +> Most document AI tools store parsed content as flat chunks in a vector DB. +> Docling-Studio v0.5 introduces a graph-native storage layer on top of Neo4j, +> preserving the full hierarchical structure of documents as first-class citizens. +> This unlocks hybrid retrieval, agentic navigation, and structural debugging — +> impossible with chunk-only stores. + +### Out of scope for v0.5.0 (roadmap mention only) +- EnrichmentWriter (entities / summaries / keywords via docling-agent) — v0.6.0 +- Agent reasoning trace viewer — v0.6.0 +- RAG hybrid (graph traversal + vector) — v0.7.0 +- Document versioning — v0.7.0+ + +--- + +## 2. Architectural principles + +### Port & adapter, with nuance + +**Write side**: one `Writer` port, **composable stages** (not alternative adapters). +Pipelines A and B are additive, not exclusive. + +``` + CORE (always) Pipeline A (RAG) Pipeline B (agent-ready, v0.6+) + ┌─────────────┐ ┌────────────────┐ ┌───────────────────┐ + │ TreeWriter │ ─────▶ │ ChunkWriter │ │ EnrichmentWriter │ + │ │ │ (existing │ │ (via docling- │ + │ │ │ OpenSearch + │ │ agent, v0.6+) │ + │ │ │ adds chunks │ │ │ + │ │ │ to Neo4j) │ │ │ + └─────────────┘ └────────────────┘ └───────────────────┘ +``` + +```python +# docling_studio/ingestion/pipeline.py +class Writer(Protocol): + def write(self, doc: DoclingDocument, ctx: IngestionContext) -> None: ... + +# Explicit composition per use case +def build_pipeline(config: PipelineConfig) -> list[Writer]: + writers = [TreeWriter(neo4j_driver)] + if config.rag_enabled: + writers.append(ChunkWriter(neo4j_driver, chunker, embedder, opensearch)) + if config.enrichment_enabled: # v0.6.0+ + writers.append(EnrichmentWriter(neo4j_driver, docling_agent)) + return writers +``` + +**Read side**: two distinct ports (same Neo4j backend, different queries). + +```python +class RAGRetrievalPort(Protocol): + def search(self, query: str, k: int) -> list[Chunk]: ... + def similar(self, chunk_id: str, k: int) -> list[Chunk]: ... + +class TreeNavigationPort(Protocol): # v0.6.0+ + def get_outline(self, doc_id: str) -> Tree: ... + def read_node(self, ref: str) -> Element: ... + def list_children(self, ref: str) -> list[Element]: ... + def walk(self, ref: str, depth: int) -> SubTree: ... +``` + +--- + +## 3. Neo4j schema + +### Constraints & indexes (created at boot) + +```cypher +// Uniqueness +CREATE CONSTRAINT document_id IF NOT EXISTS + FOR (d:Document) REQUIRE d.id IS UNIQUE; + +CREATE CONSTRAINT element_composite IF NOT EXISTS + FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE; + +CREATE CONSTRAINT page_composite IF NOT EXISTS + FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE; + +CREATE CONSTRAINT chunk_id IF NOT EXISTS + FOR (c:Chunk) REQUIRE c.id IS UNIQUE; + +// Full-text index (element text search) +CREATE FULLTEXT INDEX element_text IF NOT EXISTS + FOR (e:Element) ON EACH [e.text]; + +// Simple indexes for per-doc queries +CREATE INDEX element_doc IF NOT EXISTS FOR (e:Element) ON (e.doc_id); +CREATE INDEX chunk_doc IF NOT EXISTS FOR (c:Chunk) ON (c.doc_id); +``` + +### Data model + +```cypher +// Root document +(:Document { + id: string, // UUID or PDF hash + title: string, + source_uri: string, // path or S3 + ingested_at: datetime, + docling_version: string, + stages_applied: list, // ["tree", "chunks"] etc. + last_tree_write: datetime, + last_chunk_write: datetime, + tenant_id: string // simple multi-tenancy +}) + +// All tree elements (shared :Element label + specific label) +(:Element:SectionHeader {doc_id, self_ref, text, level, prov_page, prov_bbox}) +(:Element:Paragraph {doc_id, self_ref, text, prov_page, prov_bbox}) +(:Element:Table {doc_id, self_ref, caption, cells_json, prov_page, prov_bbox}) +(:Element:Figure {doc_id, self_ref, caption, image_uri, prov_page, prov_bbox}) +(:Element:ListItem {doc_id, self_ref, text, marker, prov_page, prov_bbox}) +(:Element:Formula {doc_id, self_ref, latex, text, prov_page, prov_bbox}) + +// Page for layout provenance +(:Page {doc_id, page_no, width, height}) + +// Chunks (Pipeline A) +(:Chunk { + id, doc_id, + text, + chunk_index, + embedding_ref, // id in OpenSearch (no inline duplication) + token_count +}) +``` + +### Relations + +```cypher +// Hierarchical structure +(:Document)-[:HAS_ROOT]->(:Element) +(:Element)-[:PARENT_OF {order: int}]->(:Element) // order preserves sequence +(:Element)-[:NEXT]->(:Element) // DFS pre-order reading + +// Layout +(:Element)-[:ON_PAGE]->(:Page) + +// Pipeline A (chunking) +(:Document)-[:HAS_CHUNK]->(:Chunk) +(:Chunk)-[:DERIVED_FROM]->(:Element) // back-reference; a chunk can span multiple elements +``` + +### Decisions + +| Decision | Choice | Rationale | +|----------|-------|---------------| +| Element composite key | `(doc_id, self_ref)` | self_ref not unique across docs | +| Multi-tenancy | `tenant_id` property on Document | Simple, filterable, migrable to multi-db later | +| Table cells | `cells_json` property | v0.5 KISS. May model `(Table)-[:HAS_CELL]->(Cell)` in v0.6+ | +| Reading order | `[:NEXT]` chain + `{order}` on `PARENT_OF` | Both views useful | +| Versioning | None (replace strategy on re-upload) | v0.5 KISS | +| APOC | Not required | Pure Cypher is sufficient for v0.5 | + +### Re-ingestion strategy + +```cypher +// Before ingesting, wipe existing +MATCH (d:Document {id: $doc_id}) +OPTIONAL MATCH (d)-[:HAS_ROOT|HAS_CHUNK]->() +DETACH DELETE d +// Then re-walk cleanly +``` + +--- + +## 4. Implementation plan (3 days) + +### Day 1 — Infra + schema + +- [ ] Add `neo4j` service to `docker-compose.yml` (`neo4j:5.15-community`, persistent volume, healthcheck) +- [ ] Add env vars (`NEO4J_URI`, `NEO4J_USER`, `NEO4J_PASSWORD`) to `.env.example` +- [ ] Create module `docling_studio/storage/neo4j/`: + - `driver.py` — neo4j-python driver wrapper (connection pool, context manager) + - `schema.py` — idempotent `bootstrap_schema()` (CREATE CONSTRAINT / INDEX at startup) + - `__init__.py` with exports +- [ ] Hook `bootstrap_schema()` in FastAPI startup +- [ ] Basic integration tests: + - Driver connection + - Schema bootstrap (idempotence verified) + - Simple round-trip: write Document, read Document, delete Document + +**Deliverable:** docker compose boots with healthy Neo4j, schema in place at init. + +### Day 2 — TreeWriter (write + read) + +- [ ] `storage/neo4j/tree_writer.py` — `DoclingDocument → Neo4j` walker + - `write_document(doc, tenant_id, driver)` in a transaction + - DFS pre-order, batched `MERGE` for perf + - Pages first, then Elements, then `PARENT_OF` / `NEXT` / `ON_PAGE` relations + - Dynamic labels based on `node.label` (`SectionHeader`, `Paragraph`, …) +- [ ] `storage/neo4j/tree_reader.py` — inverse walker `Neo4j → DoclingDocument` + - `read_document(doc_id, driver) -> DoclingDocument` + - Loads all Elements + Pages, rebuilds the Pydantic structure + - Prerequisite for v0.6 (feeding docling-agent from Neo4j) +- [ ] Integrate into existing ingestion pipeline: + - Add TreeWriter as first stage of `IngestionPipeline` + - `neo4j_enabled: bool` config toggle +- [ ] Round-trip tests: + - 3–4 varied PDFs (academic, invoice, report) + - Assertion: `doc_original == read_document(write_document(doc_original))` + - Beware dates, bbox floats (tolerance) + +**Deliverable:** A PDF uploaded to Docling-Studio is fully present in Neo4j and rebuildable. + +### Day 3 — UI + ChunkWriter + packaging + +- [ ] `storage/neo4j/chunk_writer.py`: + - After existing chunking, push each Chunk to Neo4j + - Create `(:Chunk)-[:DERIVED_FROM]->(:Element)` via source element `self_ref` + - Do NOT duplicate embeddings (stay in OpenSearch, keep `embedding_ref`) +- [ ] Frontend: new "Graph view" tab in debug panel + - Vue component with `cytoscape` (lighter, better layout API) + - FastAPI endpoint `/api/documents/{doc_id}/graph` returns nodes + edges around a scope (whole doc or subtree) + - View: vertical tree, colors per node type, click-to-zoom, hover details +- [ ] Per-document "Graph-ready" / "RAG-ready" badge in list +- [ ] README update: + - "Graph storage with Neo4j" section + - Schema diagram (Mermaid or image) + - 2–3 Cypher examples like "find all paragraphs under section X that mention Y" + - Neo4j badge in features list +- [ ] (bonus if time) "Query explorer" dev tab for live demo: Cypher editor + results + +**Deliverable:** release 0.5.0 with Neo4j visible, functional, documented. + +--- + +## 5. Proposed code structure + +``` +docling_studio/ +├── storage/ +│ ├── neo4j/ +│ │ ├── __init__.py +│ │ ├── driver.py # connection management +│ │ ├── schema.py # bootstrap_schema() +│ │ ├── tree_writer.py # DoclingDocument -> Neo4j +│ │ ├── tree_reader.py # Neo4j -> DoclingDocument +│ │ ├── chunk_writer.py # Chunks -> Neo4j +│ │ └── queries.py # shared Cypher queries +│ ├── opensearch/ # (existing) +│ └── ports.py # Writer, RAGRetrievalPort protocols +├── ingestion/ +│ └── pipeline.py # IngestionPipeline composing Writers +├── api/ +│ └── graph.py # /api/documents/{id}/graph +└── frontend/ + └── components/ + └── GraphView.vue # cytoscape + graph API fetch +``` + +--- + +## 6. Docker compose (added excerpt) + +```yaml +services: + neo4j: + image: neo4j:5.15-community + environment: + NEO4J_AUTH: ${NEO4J_USER:-neo4j}/${NEO4J_PASSWORD:-changeme} + NEO4J_PLUGINS: '["apoc"]' + NEO4J_server_memory_heap_initial__size: 512m + NEO4J_server_memory_heap_max__size: 1g + ports: + - "7474:7474" # Browser UI (demo) + - "7687:7687" # Bolt protocol + volumes: + - neo4j_data:/data + - neo4j_logs:/logs + healthcheck: + test: ["CMD-SHELL", "cypher-shell -u neo4j -p $${NEO4J_PASSWORD:-changeme} 'RETURN 1' || exit 1"] + interval: 10s + timeout: 5s + retries: 10 + + docling-studio-backend: + depends_on: + neo4j: + condition: service_healthy + environment: + NEO4J_URI: bolt://neo4j:7687 + NEO4J_USER: neo4j + NEO4J_PASSWORD: ${NEO4J_PASSWORD:-changeme} + +volumes: + neo4j_data: + neo4j_logs: +``` + +--- + +## 7. Tests + +### Unit tests +- `tests/storage/neo4j/test_schema.py` — bootstrap is idempotent +- `tests/storage/neo4j/test_tree_writer.py` — round-trip on synthetic DoclingDocument +- `tests/storage/neo4j/test_chunk_writer.py` — chunks written with correct `DERIVED_FROM` + +### Integration tests +- `tests/integration/test_ingestion_pipeline.py` — full pipeline on a real PDF +- PDF fixtures: 1 academic (complex heading hierarchy), 1 invoice (tables), 1 report (lists) + +### E2E (bonus) +- Upload PDF via UI → check structure in Neo4j Browser + +--- + +## 8. Open decisions to settle before coding + +1. **Neo4j edition**: Community (free) or AuraDB (managed) ? + - Rec: Community in Docker for v0.5.0 dev/demo. AuraDB mentioned as prod option. + +2. **Chunks: duplicate embeddings in Neo4j or OpenSearch ref ?** + - Rec: OpenSearch ref (avoid duplication; OpenSearch remains source of truth for vectors). In v0.6+, consider native Neo4j vector index. + +3. **Graph view UI: cytoscape or vis-network ?** + - Rec: **cytoscape** — lighter, better layout API, used by Neo4j itself. + +4. **Graph endpoint: return full doc or paginate ?** + - Rec: full doc for v0.5 (reasonable cap at 200 pages). Pagination in v0.6. + +5. **Error strategy**: if Neo4j is down at ingestion, fail or degrade gracefully ? + - Rec: **fail fast** for v0.5 (avoid silent inconsistencies). `neo4j_required: bool` config option. + +--- + +## 9. Hooks for later (v0.6.0+ — don't implement but prepare) + +**EnrichmentWriter (v0.6)** — will need: +- The reader (Neo4j → DoclingDocument) to re-materialize the doc, feed docling-agent, re-patch enrichments +- A stage addable to `IngestionPipeline` without touching other stages +- An `:Entity` label (not created in v0.5 but schema-compatible) + +**Agent reasoning trace viewer (v0.6)** — will need: +- An event stream (WebSocket) that v0.5 already prepares via the reactive UI +- A node_ref ↔ Element correlation in Neo4j (our composite `self_ref` key is enough) + +**TreeNavigationPort (v0.7)** — will need: +- Optimized Cypher queries for descendant/ancestor walk (indexes already provisioned) + +--- + +## 10. v0.5.0 success criteria + +**Must have:** +- [ ] A PDF uploaded to Docling-Studio is in Neo4j with structure preserved +- [ ] Neo4j Browser shows the graph and is manually explorable +- [ ] A graph visual in the Docling-Studio UI works +- [ ] `docker compose up` works zero-config +- [ ] README mentions Neo4j and describes the schema + +**Nice to have (decreasing priority):** +- [ ] Graph-ready / RAG-ready badge per doc +- [ ] Live query explorer in the UI +- [ ] 2–3 example queries in README that do something impossible with vector-only + +**For the hackathon (post-release):** +- [ ] 60s video: upload PDF → structure in Neo4j → cross-doc query impossible in vector-only +- [ ] HackerNoon post explaining "graph-native documents" positioning +- [ ] Explicit Neo4j partnership mention + +--- + +## 11. Fundamental architectural decisions recap + +| Question | Answer | +|----------|---------| +| Is Neo4j source of truth or cache ? | **Source of truth** for structure. OpenSearch remains source of truth for embeddings. | +| Does chunking go away ? | No, v0.5.0 keeps existing chunking. "Chunkless" is Pipeline B, v0.6+. | +| Can it be toggled per doc ? | Yes — `stages_applied` on Document + pipeline config | +| What about OpenSearch ? | Stays, stores vectors. Neo4j tracks `(:Chunk)-[:DERIVED_FROM]->(:Element)` links. | +| Multi-tenancy ? | `tenant_id` property on Document, Cypher filter | +| Versioning ? | None for v0.5.0 — replace strategy on re-upload | + +--- + +## Appendix — Demo queries + +### Query 1 — All "Methods" sections across documents +```cypher +MATCH (d:Document)-[:HAS_ROOT]->(:Element)-[:PARENT_OF*]->(s:SectionHeader) +WHERE toLower(s.text) CONTAINS 'method' +RETURN d.title, s.text, s.level +``` + +### Query 2 — Context of a chunk (parent + siblings) +```cypher +MATCH (c:Chunk {id: $chunk_id})-[:DERIVED_FROM]->(e:Element) +MATCH (e)<-[:PARENT_OF]-(parent:Element) +MATCH (parent)-[:PARENT_OF]->(sibling:Element) +RETURN parent, collect(sibling) AS siblings +``` + +### Query 3 — All tables from a document type +```cypher +MATCH (d:Document)-[:HAS_ROOT]->(:Element)-[:PARENT_OF*]->(t:Table) +WHERE d.source_uri CONTAINS 'invoices/' +RETURN d.title, t.caption, t.cells_json +``` + +### Query 4 — Direct children of a section (ordered) +```cypher +MATCH (s:Element {doc_id: $doc_id, self_ref: $section_ref}) +MATCH (s)-[pof:PARENT_OF]->(child) +RETURN child +ORDER BY pof.order +``` + +--- + +*Single reference doc for Neo4j v0.5.0 implementation. +Read this first in the implementation thread.* diff --git a/document-parser/infra/neo4j/__init__.py b/document-parser/infra/neo4j/__init__.py new file mode 100644 index 0000000..23be98d --- /dev/null +++ b/document-parser/infra/neo4j/__init__.py @@ -0,0 +1,10 @@ +"""Neo4j storage adapter — graph-native document structure. + +Provides a thin driver wrapper, idempotent schema bootstrap, and +walkers between DoclingDocument and the graph model. +""" + +from infra.neo4j.driver import Neo4jDriver, close_driver, get_driver +from infra.neo4j.schema import bootstrap_schema + +__all__ = ["Neo4jDriver", "bootstrap_schema", "close_driver", "get_driver"] diff --git a/document-parser/infra/neo4j/driver.py b/document-parser/infra/neo4j/driver.py new file mode 100644 index 0000000..0745d59 --- /dev/null +++ b/document-parser/infra/neo4j/driver.py @@ -0,0 +1,48 @@ +"""Async Neo4j driver wrapper. + +Owns a single `AsyncDriver` per process. Callers acquire it via +`get_driver()` and must call `close_driver()` at shutdown. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +from neo4j import AsyncDriver, AsyncGraphDatabase + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class Neo4jDriver: + driver: AsyncDriver + database: str = "neo4j" + + +_instance: Neo4jDriver | None = None + + +async def get_driver(uri: str, user: str, password: str, database: str = "neo4j") -> Neo4jDriver: + """Return the process-wide driver, creating it on first call. + + Verifies connectivity once at creation — raises if the server is unreachable. + """ + global _instance + if _instance is not None: + return _instance + + driver = AsyncGraphDatabase.driver(uri, auth=(user, password)) + await driver.verify_connectivity() + logger.info("Neo4j driver connected to %s (db=%s)", uri, database) + _instance = Neo4jDriver(driver=driver, database=database) + return _instance + + +async def close_driver() -> None: + global _instance + if _instance is None: + return + await _instance.driver.close() + _instance = None + logger.info("Neo4j driver closed") diff --git a/document-parser/infra/neo4j/schema.py b/document-parser/infra/neo4j/schema.py new file mode 100644 index 0000000..9307d3b --- /dev/null +++ b/document-parser/infra/neo4j/schema.py @@ -0,0 +1,51 @@ +"""Idempotent Neo4j schema bootstrap. + +Runs at backend startup. All statements use `IF NOT EXISTS`, so calling +this multiple times is safe — it's the contract integration tests rely on. +""" + +from __future__ import annotations + +import logging + +from infra.neo4j.driver import Neo4jDriver + +logger = logging.getLogger(__name__) + + +CONSTRAINTS: tuple[str, ...] = ( + "CREATE CONSTRAINT document_id IF NOT EXISTS " + "FOR (d:Document) REQUIRE d.id IS UNIQUE", + "CREATE CONSTRAINT element_composite IF NOT EXISTS " + "FOR (e:Element) REQUIRE (e.doc_id, e.self_ref) IS UNIQUE", + "CREATE CONSTRAINT page_composite IF NOT EXISTS " + "FOR (p:Page) REQUIRE (p.doc_id, p.page_no) IS UNIQUE", + "CREATE CONSTRAINT chunk_id IF NOT EXISTS " + "FOR (c:Chunk) REQUIRE c.id IS UNIQUE", +) + +INDEXES: tuple[str, ...] = ( + "CREATE INDEX element_doc IF NOT EXISTS FOR (e:Element) ON (e.doc_id)", + "CREATE INDEX chunk_doc IF NOT EXISTS FOR (c:Chunk) ON (c.doc_id)", +) + +FULLTEXT_INDEXES: tuple[str, ...] = ( + "CREATE FULLTEXT INDEX element_text IF NOT EXISTS " + "FOR (e:Element) ON EACH [e.text]", +) + + +async def bootstrap_schema(neo: Neo4jDriver) -> None: + """Create constraints and indexes required by the graph model. + + Idempotent: safe to call on every startup. + """ + async with neo.driver.session(database=neo.database) as session: + for stmt in (*CONSTRAINTS, *INDEXES, *FULLTEXT_INDEXES): + await session.run(stmt) + logger.info( + "Neo4j schema bootstrapped (%d constraints, %d indexes, %d fulltext)", + len(CONSTRAINTS), + len(INDEXES), + len(FULLTEXT_INDEXES), + ) diff --git a/document-parser/infra/settings.py b/document-parser/infra/settings.py index cc322fe..b114045 100644 --- a/document-parser/infra/settings.py +++ b/document-parser/infra/settings.py @@ -25,6 +25,9 @@ class Settings: batch_page_size: int = 0 # 0 = disabled, > 0 = pages per batch opensearch_url: str = "" # empty = disabled embedding_url: str = "" # empty = disabled (e.g. http://localhost:8001) + neo4j_uri: str = "" # empty = disabled (e.g. bolt://neo4j:7687) + neo4j_user: str = "neo4j" + neo4j_password: str = "changeme" opensearch_default_limit: int = 1000 # max chunks returned by get_chunks embedding_dimension: int = 384 # Granite Embedding 30M / all-MiniLM-L6-v2 upload_dir: str = "./uploads" @@ -102,6 +105,9 @@ class Settings: batch_page_size=int(os.environ.get("BATCH_PAGE_SIZE", "10")), opensearch_url=os.environ.get("OPENSEARCH_URL", ""), embedding_url=os.environ.get("EMBEDDING_URL", ""), + neo4j_uri=os.environ.get("NEO4J_URI", ""), + neo4j_user=os.environ.get("NEO4J_USER", "neo4j"), + neo4j_password=os.environ.get("NEO4J_PASSWORD", "changeme"), opensearch_default_limit=int(os.environ.get("OPENSEARCH_DEFAULT_LIMIT", "1000")), embedding_dimension=int(os.environ.get("EMBEDDING_DIMENSION", "384")), upload_dir=os.environ.get("UPLOAD_DIR", "./uploads"), diff --git a/document-parser/main.py b/document-parser/main.py index a163e85..1c41921 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -93,6 +93,28 @@ def _build_analysis_service( ) +async def _init_neo4j(): + """Initialize the Neo4j driver and bootstrap schema — skip if not configured.""" + if not settings.neo4j_uri: + logger.info("Neo4j disabled (NEO4J_URI not set)") + return None + + from infra.neo4j import Neo4jDriver, bootstrap_schema, get_driver + + try: + neo = await get_driver( + settings.neo4j_uri, + settings.neo4j_user, + settings.neo4j_password, + ) + await bootstrap_schema(neo) + logger.info("Neo4j ready (uri=%s)", settings.neo4j_uri) + return neo + except Exception: + logger.exception("Neo4j init failed — continuing without graph storage") + return None + + def _build_ingestion_service() -> IngestionService | None: """Build the ingestion service — only if embedding + opensearch are configured.""" if not settings.embedding_url or not settings.opensearch_url: @@ -150,8 +172,15 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: if ingestion_service is not None: app.include_router(ingestion_router) logger.info("Ingestion router mounted") + app.state.neo4j = await _init_neo4j() logger.info("Docling Studio backend ready (engine=%s)", settings.conversion_engine) - yield + try: + yield + finally: + if app.state.neo4j is not None: + from infra.neo4j import close_driver + + await close_driver() app = FastAPI( diff --git a/document-parser/requirements.txt b/document-parser/requirements.txt index 2319b46..1137357 100644 --- a/document-parser/requirements.txt +++ b/document-parser/requirements.txt @@ -8,3 +8,4 @@ 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 +neo4j>=5.15.0,<6.0.0 diff --git a/document-parser/tests/neo4j/__init__.py b/document-parser/tests/neo4j/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/document-parser/tests/neo4j/conftest.py b/document-parser/tests/neo4j/conftest.py new file mode 100644 index 0000000..9c77d44 --- /dev/null +++ b/document-parser/tests/neo4j/conftest.py @@ -0,0 +1,36 @@ +"""Shared fixtures for Neo4j integration tests. + +These tests are skipped unless a live Neo4j is reachable via NEO4J_TEST_URI +(defaulting to bolt://localhost:7687). CI spins up `neo4j:5.15-community` +alongside the job; locally run `docker compose -f docker-compose.dev.yml up neo4j`. +""" + +from __future__ import annotations + +import os + +import pytest + +from infra.neo4j import close_driver, get_driver + + +def _cfg() -> tuple[str, str, str]: + return ( + os.environ.get("NEO4J_TEST_URI", "bolt://localhost:7687"), + os.environ.get("NEO4J_TEST_USER", "neo4j"), + os.environ.get("NEO4J_TEST_PASSWORD", "changeme"), + ) + + +@pytest.fixture +async def neo4j_driver(): + uri, user, password = _cfg() + try: + neo = await get_driver(uri, user, password) + except Exception as exc: + pytest.skip(f"Neo4j not reachable at {uri}: {exc}") + # Wipe DB before each test — integration tests assume an empty graph. + async with neo.driver.session(database=neo.database) as session: + await session.run("MATCH (n) DETACH DELETE n") + yield neo + await close_driver() diff --git a/document-parser/tests/neo4j/test_document_roundtrip.py b/document-parser/tests/neo4j/test_document_roundtrip.py new file mode 100644 index 0000000..1303d85 --- /dev/null +++ b/document-parser/tests/neo4j/test_document_roundtrip.py @@ -0,0 +1,32 @@ +"""Minimal Document node round-trip — validates the driver + schema end-to-end.""" + +from __future__ import annotations + +from infra.neo4j.schema import bootstrap_schema + + +async def test_document_write_read_delete(neo4j_driver): + await bootstrap_schema(neo4j_driver) + + async with neo4j_driver.driver.session(database=neo4j_driver.database) as session: + await session.run( + "CREATE (d:Document {id: $id, title: $title, tenant_id: $tenant})", + id="doc-42", + title="Round-trip fixture", + tenant="default", + ) + + result = await session.run( + "MATCH (d:Document {id: $id}) RETURN d.title AS title, d.tenant_id AS tenant", + id="doc-42", + ) + record = await result.single() + assert record is not None + assert record["title"] == "Round-trip fixture" + assert record["tenant"] == "default" + + await session.run("MATCH (d:Document {id: $id}) DETACH DELETE d", id="doc-42") + gone = await ( + await session.run("MATCH (d:Document {id: $id}) RETURN d", id="doc-42") + ).single() + assert gone is None diff --git a/document-parser/tests/neo4j/test_driver.py b/document-parser/tests/neo4j/test_driver.py new file mode 100644 index 0000000..abb42ea --- /dev/null +++ b/document-parser/tests/neo4j/test_driver.py @@ -0,0 +1,10 @@ +"""Neo4j driver connectivity smoke test.""" + +from __future__ import annotations + + +async def test_driver_connects_and_runs_cypher(neo4j_driver): + async with neo4j_driver.driver.session(database=neo4j_driver.database) as session: + result = await session.run("RETURN 1 AS x") + record = await result.single() + assert record["x"] == 1 diff --git a/document-parser/tests/neo4j/test_schema.py b/document-parser/tests/neo4j/test_schema.py new file mode 100644 index 0000000..4e6e2c2 --- /dev/null +++ b/document-parser/tests/neo4j/test_schema.py @@ -0,0 +1,38 @@ +"""Schema bootstrap is idempotent and produces the expected constraints/indexes.""" + +from __future__ import annotations + +from infra.neo4j.schema import CONSTRAINTS, FULLTEXT_INDEXES, INDEXES, bootstrap_schema + + +async def _count_schema(neo4j_driver) -> tuple[int, int]: + async with neo4j_driver.driver.session(database=neo4j_driver.database) as session: + constraints = await (await session.run("SHOW CONSTRAINTS")).data() + indexes = await (await session.run("SHOW INDEXES")).data() + return len(constraints), len(indexes) + + +async def test_bootstrap_is_idempotent(neo4j_driver): + await bootstrap_schema(neo4j_driver) + first = await _count_schema(neo4j_driver) + + # Running a second time must not duplicate anything. + await bootstrap_schema(neo4j_driver) + second = await _count_schema(neo4j_driver) + + assert first == second + # Sanity: we created at least what we declared. + assert first[0] >= len(CONSTRAINTS) + assert first[1] >= len(INDEXES) + len(FULLTEXT_INDEXES) + + +async def test_document_id_is_unique(neo4j_driver): + await bootstrap_schema(neo4j_driver) + async with neo4j_driver.driver.session(database=neo4j_driver.database) as session: + await session.run("CREATE (d:Document {id: 'doc-1', title: 'first'})") + with_err: Exception | None = None + try: + await session.run("CREATE (d:Document {id: 'doc-1', title: 'dup'})") + except Exception as exc: + with_err = exc + assert with_err is not None, "unique constraint on Document.id must reject duplicates"