ragleap-graph 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,32 @@
1
+ # Environment
2
+ .env
3
+ *.env.local
4
+
5
+ # Python
6
+ __pycache__/
7
+ *.py[cod]
8
+ *.egg-info/
9
+ venv/
10
+ .venv/
11
+
12
+ # Node
13
+ node_modules/
14
+
15
+ # Docker
16
+ docker-compose.override.yml
17
+
18
+ # IDE
19
+ .vscode/
20
+ .idea/
21
+
22
+ # OS
23
+ .DS_Store
24
+ Thumbs.db
25
+
26
+ # Logs
27
+ *.log
28
+ .env
29
+
30
+ # Package build artifacts
31
+ dist/
32
+ build/
@@ -0,0 +1,24 @@
1
+ # Changelog
2
+
3
+ All notable changes to `ragleap-graph` are documented here. Format loosely
4
+ follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
5
+
6
+ ## [Unreleased]
7
+
8
+ ## [0.1.0]
9
+
10
+ ### Added
11
+
12
+ - Initial implementation: `GraphConfig` and `GraphIndex`, knowledge-graph-augmented retrieval for RAG systems via Neo4j.
13
+ - Ported from a real, production `GraphService` (Neo4j-backed, regex entity extraction, co-occurrence graph construction), with deliberate adaptations for standalone open-source use:
14
+ - `workspace_id` (required, Django multi-tenant SaaS concept) generalized to an optional `namespace=` parameter, defaulting to `None` — mirrors `ragleap-rag`'s own `metadata_filter=` convention.
15
+ - Django `settings` coupling replaced with an explicit `GraphConfig` dataclass — zero framework dependency, same pattern as `ragleap-rag`'s `EmbeddingConfig`/`ProviderConfig`.
16
+ - A hardcoded, deployment-specific domain vocabulary list (radiation-safety/medical-physics terms) found in the source was **not** carried over as a hardcoded default — it's now an optional `domain_terms=` parameter instead, since shipping someone else's internal business vocabulary as a public library default would be inappropriate and actively wrong for other domains.
17
+ - 5 core methods: `upsert_document`, `find_documents_by_entities`, `search_related_entities`, `document_entities`, `health_check`.
18
+ - **Real security fix vs. the source this was ported from**: `search_related_entities()`'s `max_depth` parameter is validated as a bounded integer (1-10) before being used in Cypher query construction. The source interpolated this value into the query string unvalidated (`.format(max_depth=max_depth)`) — necessary because Neo4j's Cypher can't parameterize variable-length path bounds normally, but that means unvalidated input would be a real injection vector in a public library API accepting arbitrary caller input.
19
+ - **Real bug found and fixed during implementation** (not just written and assumed correct): `GraphDatabase.driver()` is lazy — it does not attempt a real connection at construction time. Without an explicit `verify_connectivity()` call, a bad URI/credentials would silently leave `self.driver` set to a non-functional object instead of `None`, breaking the documented "driver is None means unavailable" contract. Caught by a real test failure, not assumed correct — see the test suite for the exact scenario.
20
+ - 26 tests: pure-logic tests (entity normalization, extraction, `max_depth` validation) requiring no live Neo4j, driver-unavailable graceful-degradation tests for every public method, and one live end-to-end integration test (write, query three ways, verified self-cleanup — 0 nodes left behind) that runs automatically when real `NEO4J_URI`/`NEO4J_USER`/`NEO4J_PASSWORD` credentials are present, and skips cleanly otherwise.
21
+ - CI: dedicated `ragleap-graph-tests` job with a real Neo4j service container, mirroring `ragleap-rag-tests`' pattern — confirmed passing in GitHub's own CI environment, not just locally.
22
+ - **This is the first package in the ragleap ecosystem with genuine, complete live verification from day one** — unlike several `ragleap-rag` vector backends, which shipped honestly labeled "code-complete but unverified" due to no account access. Real Neo4j access was already available, so every method here has actually been proven to work end-to-end, not just tested in isolation.
23
+
24
+ - **Second real bug found and fixed during implementation**: `search_related_entities()` originally validated `max_depth` *after* checking driver availability, so invalid `max_depth` input silently returned `[]` instead of raising — when it should raise regardless of connection state. Caught by a real test failure. Fixed by validating `max_depth` first, before any driver-availability check.
@@ -0,0 +1,159 @@
1
+ # ragleap-graph — Build Handoff Document
2
+
3
+ **Purpose:** continuity doc for a new chat session to resume building
4
+ `ragleap-graph`. Scoping is done; implementation hasn't started.
5
+
6
+ ---
7
+
8
+ ## Current state
9
+
10
+ - **Package skeleton created** at `~/ragleap-core/packages/ragleap-graph/`
11
+ on the VPS (`srv1477778`)
12
+ - `pyproject.toml` written and matches `ragleap-rag`'s exact conventions
13
+ (hatchling build backend, same author/license/classifiers pattern).
14
+ Verified complete via direct `cat` inspection (not `tomllib` — that
15
+ module wasn't available in the system Python being used; real
16
+ validation will happen when `hatchling` actually builds the package)
17
+ - Directory structure: `src/ragleap_graph/` and `tests/` both exist,
18
+ currently empty
19
+ - **Nothing has been coded yet** — no `GraphConfig`, no `GraphIndex`,
20
+ no tests
21
+
22
+ ## Where the real production code lives (read this first, don't guess)
23
+
24
+ `cdprod` → `/var/www/ragleap/backends/claude_backend` (the private
25
+ production Django backend — NOT the `ragleap-core` public monorepo,
26
+ don't confuse the two, this has bitten this session before).
27
+
28
+ Core files, already reviewed this session:
29
+ - `retrieval/graph_service.py` — 563 lines, the real `GraphService` class
30
+ - `retrieval/neo4j_startup.py` — 293 lines, connection/startup logic
31
+
32
+ 28 files across 8 Django apps reference Neo4j/graph in some way, but
33
+ match-density analysis (`grep -c`) confirmed these two files are where
34
+ the actual logic lives — everything else is either a consumer calling
35
+ into `GraphService` or a passing reference (settings flag, migration
36
+ field). Full file list and match counts are in this session's history
37
+ if needed again, but the two files above are the ones that matter.
38
+
39
+ ## The 5 real methods to port (verified via direct code review)
40
+
41
+ | Production method (Django-coupled) | Target public API | What it does |
42
+ |---|---|---|
43
+ | `upsert_document_graph(document_id, workspace_id, document_title, chunks, language=None, max_entities=80, max_pairs=150)` | `graph.upsert_document(...)` | Regex-based entity extraction per chunk, builds co-occurrence edges between entities in the same chunk, `MERGE`-based idempotent Neo4j writes |
44
+ | `find_documents_by_entities(entity_names, workspace_id, limit=25)` | `graph.find_documents_by_entities(...)` | Given entity names, returns documents connected to them, ranked by matched-entity count + summed edge weight |
45
+ | `search_related_entities(entity_names, workspace_id, max_depth=2, limit=10)` | `graph.search_related_entities(...)` | Graph traversal — finds entities related via N-hop paths |
46
+ | `find_document_entities(document_id, workspace_id, entity_types=None)` | `graph.document_entities(...)` | Per-document entity lookup, optional type filter |
47
+ | `health_check()` | `graph.health_check()` | Already generic, minimal changes needed |
48
+
49
+ Full method bodies (Cypher queries, exact logic) were reviewed line by
50
+ line this session — if this handoff doc is being read fresh without
51
+ that history, re-fetch them from the production file directly rather
52
+ than guessing at the Cypher:
53
+
54
+ ```bash
55
+ cdprod
56
+ sed -n '1,60p' ./retrieval/graph_service.py # class init, entity normalization
57
+ sed -n '163,326p' ./retrieval/graph_service.py # upsert_document_graph, full body
58
+ sed -n '326,548p' ./retrieval/graph_service.py # find_documents_by_entities, search_related_entities, find_document_entities, health_check
59
+ ```
60
+
61
+ ## Two required adaptations — decided, not open questions
62
+
63
+ ### 1. `workspace_id` → optional `namespace=`
64
+
65
+ **Decision made and confirmed this session**: production code requires
66
+ `workspace_id` on every method (Django multi-tenant SaaS assumption).
67
+ For the open-source package, this becomes an **optional** `namespace=`
68
+ parameter, defaulting to `None`:
69
+
70
+ ```python
71
+ graph.upsert_document(document_id="...", title="...", chunks=[...]) # single global graph
72
+ graph.upsert_document(document_id="...", title="...", chunks=[...], namespace="acme") # multi-tenant isolated
73
+ ```
74
+
75
+ This mirrors `ragleap-rag`'s own `metadata_filter=` pattern — optional,
76
+ generic, doesn't force one usage model onto every user. This is the
77
+ established, deliberate convention for this ecosystem; don't relitigate
78
+ it without a real reason.
79
+
80
+ ### 2. Django `settings` coupling → explicit `GraphConfig`
81
+
82
+ Production code does `from django.conf import settings` and
83
+ `getattr(settings, 'NEO4J_URI', ...)`. The package needs a config
84
+ dataclass instead, matching `ragleap-rag`'s `EmbeddingConfig`/
85
+ `ProviderConfig` pattern:
86
+
87
+ ```python
88
+ from ragleap_graph import GraphIndex, GraphConfig
89
+
90
+ graph = GraphIndex(config=GraphConfig(
91
+ uri="bolt://localhost:7687",
92
+ user="neo4j",
93
+ password="...",
94
+ ))
95
+ ```
96
+
97
+ ## One real security fix required before shipping
98
+
99
+ `search_related_entities()` builds its Cypher query using
100
+ **`.format(max_depth=max_depth)`** to inject the traversal depth
101
+ directly into the query string (`[*1..{max_depth}]`), rather than using
102
+ parameterized query args like every other method in the file. This is
103
+ likely a genuine Neo4j limitation (Cypher can't parameterize
104
+ variable-length path bounds normally) — but as a **public library API**
105
+ accepting `max_depth` from any caller, this needs explicit validation
106
+ (`isinstance(max_depth, int)` and a sane upper bound, e.g. reject
107
+ values > 10) before string formatting, or it's a real Cypher injection
108
+ vector. Don't carry this over unguarded.
109
+
110
+ ## Package conventions to follow (from `ragleap-rag`, don't reinvent)
111
+
112
+ - `pyproject.toml`: hatchling build backend, same classifiers/license
113
+ pattern, `[project.optional-dependencies]` with a `test` extra
114
+ - Version bumps go in both `pyproject.toml` and
115
+ `src/ragleap_graph/__init__.py`'s `__version__` — grep-check both
116
+ - CHANGELOG.md discipline: never rewrite past entries, add corrections
117
+ as new entries
118
+ - For any third-party SDK usage (the `neo4j` Python driver), verify
119
+ exact method signatures via `inspect.signature()` against the real
120
+ installed package before writing code — don't trust memory or docs.
121
+ This caught a real pre-ship bug in `PineconeBackend` before; apply the
122
+ same discipline here.
123
+ - Full test suite must pass before any commit is considered done — this
124
+ project's standard, demonstrated repeatedly (238/238 for `ragleap-rag`)
125
+ - Live-verification advantage specific to this package: **you have real,
126
+ working Neo4j access already** (it's live in production) — meaning
127
+ `ragleap-graph` can be fully live-verified from day one, unlike several
128
+ `ragleap-rag` vector backends that shipped honestly labeled
129
+ "code-complete but unverified" due to no account access. Don't skip
130
+ the live verification just because it's more available this time —
131
+ make sure it actually happens.
132
+
133
+ ## Immediate next steps for the new session
134
+
135
+ 1. Write `src/ragleap_graph/__init__.py` with `GraphConfig` and
136
+ `GraphIndex` classes, porting the 5 methods above with the two
137
+ adaptations (namespace, config) and the `max_depth` fix applied
138
+ 2. Write a real `README.md` for the package (required — `pyproject.toml`
139
+ already references `readme = "README.md"` and will fail to build
140
+ without it)
141
+ 3. Write tests — mirror `ragleap-rag`'s test structure (constructor
142
+ validation, pure logic tests where possible without a live Neo4j,
143
+ then real integration tests against your actual working Neo4j
144
+ instance)
145
+ 4. Verify build: `python3 -m build` inside
146
+ `~/ragleap-core/packages/ragleap-graph/`
147
+ 5. Only after tests pass and a real Neo4j round-trip is verified: decide
148
+ whether to publish to PyPI or keep it in-repo longer
149
+
150
+ ## Session artifacts also produced today (separate from this task)
151
+
152
+ - `ragleap-rag` v0.11.2 published to PyPI (README diagram fixes,
153
+ benchmark DB isolation bug fixed)
154
+ - Full 9-page GitHub wiki written and pushed
155
+ (`github.com/antonyrag/ragleap-core/wiki`)
156
+ - Marketing/launch copy prepared in a separate handoff doc (Product
157
+ Hunt, LinkedIn, Dev.to) — for `ragleap-rag` specifically, not
158
+ `ragleap.com` the commercial platform (these are different products,
159
+ don't mix their messaging)
@@ -0,0 +1,61 @@
1
+ Metadata-Version: 2.4
2
+ Name: ragleap-graph
3
+ Version: 0.1.0
4
+ Summary: Knowledge-graph-augmented retrieval for RAG systems — entity extraction, co-occurrence graphs, and graph-based document retrieval via Neo4j. Framework-agnostic, optional multi-tenant namespacing.
5
+ Project-URL: Homepage, https://github.com/antonyrag/ragleap-core
6
+ Project-URL: Repository, https://github.com/antonyrag/ragleap-core
7
+ Project-URL: Issues, https://github.com/antonyrag/ragleap-core/issues
8
+ Author-email: Antony <antony@ragleap.com>
9
+ License-Expression: MIT
10
+ Keywords: entity-extraction,graph-retrieval,knowledge-graph,neo4j,rag
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: License :: OSI Approved :: MIT License
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
19
+ Requires-Python: >=3.10
20
+ Requires-Dist: neo4j>=5.0.0
21
+ Provides-Extra: test
22
+ Requires-Dist: pytest-asyncio>=0.24.0; extra == 'test'
23
+ Requires-Dist: pytest>=8.0.0; extra == 'test'
24
+ Description-Content-Type: text/markdown
25
+
26
+ # ragleap-graph
27
+
28
+ Knowledge-graph-augmented retrieval for RAG systems — entity extraction, co-occurrence graphs, and graph-based document retrieval via Neo4j.
29
+
30
+ ```bash
31
+ pip install ragleap-graph
32
+ ```
33
+
34
+ ## Quickstart
35
+
36
+ ```python
37
+ from ragleap_graph import GraphConfig, GraphIndex
38
+
39
+ graph = GraphIndex(config=GraphConfig(
40
+ uri="bolt://localhost:7687",
41
+ user="neo4j",
42
+ password="...",
43
+ ))
44
+
45
+ graph.upsert_document(
46
+ document_id="doc-1",
47
+ title="Q3 Report",
48
+ chunks=[{"text": "Acme Corp reported strong Q3 revenue growth."}],
49
+ )
50
+
51
+ docs = graph.find_documents_by_entities(["Acme Corp"])
52
+ related = graph.search_related_entities(["Acme Corp"], max_depth=2)
53
+ ```
54
+
55
+ ## Status
56
+
57
+ Early alpha (v0.1.0). Ported from a real production `GraphService`, adapted for standalone open-source use — see `HANDOFF.md` for the full design history.
58
+
59
+ ## License
60
+
61
+ MIT
@@ -0,0 +1,36 @@
1
+ # ragleap-graph
2
+
3
+ Knowledge-graph-augmented retrieval for RAG systems — entity extraction, co-occurrence graphs, and graph-based document retrieval via Neo4j.
4
+
5
+ ```bash
6
+ pip install ragleap-graph
7
+ ```
8
+
9
+ ## Quickstart
10
+
11
+ ```python
12
+ from ragleap_graph import GraphConfig, GraphIndex
13
+
14
+ graph = GraphIndex(config=GraphConfig(
15
+ uri="bolt://localhost:7687",
16
+ user="neo4j",
17
+ password="...",
18
+ ))
19
+
20
+ graph.upsert_document(
21
+ document_id="doc-1",
22
+ title="Q3 Report",
23
+ chunks=[{"text": "Acme Corp reported strong Q3 revenue growth."}],
24
+ )
25
+
26
+ docs = graph.find_documents_by_entities(["Acme Corp"])
27
+ related = graph.search_related_entities(["Acme Corp"], max_depth=2)
28
+ ```
29
+
30
+ ## Status
31
+
32
+ Early alpha (v0.1.0). Ported from a real production `GraphService`, adapted for standalone open-source use — see `HANDOFF.md` for the full design history.
33
+
34
+ ## License
35
+
36
+ MIT
@@ -0,0 +1,39 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "ragleap-graph"
7
+ version = "0.1.0"
8
+ description = "Knowledge-graph-augmented retrieval for RAG systems — entity extraction, co-occurrence graphs, and graph-based document retrieval via Neo4j. Framework-agnostic, optional multi-tenant namespacing."
9
+ readme = "README.md"
10
+ license = "MIT"
11
+ requires-python = ">=3.10"
12
+ authors = [
13
+ { name = "Antony", email = "antony@ragleap.com" }
14
+ ]
15
+ keywords = ["knowledge-graph", "rag", "neo4j", "entity-extraction", "graph-retrieval"]
16
+ classifiers = [
17
+ "Development Status :: 3 - Alpha",
18
+ "Intended Audience :: Developers",
19
+ "License :: OSI Approved :: MIT License",
20
+ "Programming Language :: Python :: 3",
21
+ "Programming Language :: Python :: 3.10",
22
+ "Programming Language :: Python :: 3.11",
23
+ "Programming Language :: Python :: 3.12",
24
+ "Topic :: Software Development :: Libraries :: Python Modules",
25
+ ]
26
+ dependencies = [
27
+ "neo4j>=5.0.0",
28
+ ]
29
+
30
+ [project.optional-dependencies]
31
+ test = ["pytest>=8.0.0", "pytest-asyncio>=0.24.0"]
32
+
33
+ [project.urls]
34
+ Homepage = "https://github.com/antonyrag/ragleap-core"
35
+ Repository = "https://github.com/antonyrag/ragleap-core"
36
+ Issues = "https://github.com/antonyrag/ragleap-core/issues"
37
+
38
+ [tool.hatch.build.targets.wheel]
39
+ packages = ["src/ragleap_graph"]
@@ -0,0 +1,536 @@
1
+ """
2
+ ragleap-graph: knowledge-graph-augmented retrieval for RAG systems.
3
+
4
+ Ported from a real, production GraphService (Neo4j-backed, regex entity
5
+ extraction, co-occurrence graph construction) with two deliberate
6
+ adaptations for standalone open-source use:
7
+
8
+ 1. `workspace_id` (required, Django multi-tenant SaaS concept) becomes
9
+ an optional `namespace=` parameter, defaulting to None — mirrors
10
+ ragleap-rag's own `metadata_filter=` convention.
11
+ 2. Django `settings` coupling becomes an explicit `GraphConfig`
12
+ dataclass — zero framework dependency, same pattern as ragleap-rag's
13
+ `EmbeddingConfig`/`ProviderConfig`.
14
+
15
+ A real security fix vs. the source this was ported from: `max_depth` in
16
+ `search_related_entities()` is validated as a bounded integer before
17
+ being used in Cypher query construction (the underlying Neo4j
18
+ variable-length-path syntax can't be parameterized normally, so the
19
+ depth value has to be validated before string formatting, not just
20
+ trusted).
21
+ """
22
+ from __future__ import annotations
23
+
24
+ import logging
25
+ import re
26
+ from collections import Counter
27
+ from dataclasses import dataclass
28
+ from typing import Any, Dict, List, Optional, Tuple
29
+
30
+ try:
31
+ from neo4j import GraphDatabase
32
+ except ImportError:
33
+ GraphDatabase = None
34
+
35
+ logger = logging.getLogger(__name__)
36
+
37
+ __version__ = "0.1.0"
38
+
39
+ # Hard ceiling on traversal depth — prevents both runaway queries and,
40
+ # since max_depth is string-interpolated into Cypher (see note above),
41
+ # guards against malformed/malicious input reaching the query text.
42
+ MAX_ALLOWED_DEPTH = 10
43
+
44
+
45
+ @dataclass
46
+ class GraphConfig:
47
+ """Connection config for a Neo4j-backed GraphIndex.
48
+
49
+ No framework dependency — pass values directly rather than relying
50
+ on ambient settings (Django, env vars, etc.). Callers who want
51
+ env-var-driven config can build one with `GraphConfig(uri=os.environ[...])`
52
+ themselves.
53
+ """
54
+ uri: str = "bolt://localhost:7687"
55
+ user: str = "neo4j"
56
+ password: str = ""
57
+
58
+
59
+ class GraphIndex:
60
+ """
61
+ Neo4j-backed knowledge graph index for entity-augmented retrieval.
62
+
63
+ Gracefully degrades if the `neo4j` package isn't installed or the
64
+ driver can't connect — `self.driver` will be None and all query
65
+ methods return empty results rather than raising, matching the
66
+ behavior of the production service this was ported from. Always
67
+ safe to construct; check `health_check()` if you need to confirm
68
+ connectivity before relying on it.
69
+ """
70
+
71
+ def __init__(self, config: Optional[GraphConfig] = None):
72
+ self.config = config or GraphConfig()
73
+ self.driver = None
74
+
75
+ if GraphDatabase is None:
76
+ logger.warning(
77
+ "neo4j package not installed. Install with: pip install ragleap-graph[neo4j] "
78
+ "(or just `pip install neo4j`). Graph queries will be unavailable."
79
+ )
80
+ return
81
+
82
+ try:
83
+ self.driver = GraphDatabase.driver(
84
+ self.config.uri,
85
+ auth=(self.config.user, self.config.password),
86
+ )
87
+ # GraphDatabase.driver() is lazy - it does not attempt a real
88
+ # connection until first use. Without an explicit check here,
89
+ # a bad URI/credentials would silently leave self.driver set
90
+ # to a non-functional object instead of None, breaking the
91
+ # "always safe to construct, check driver is None" contract
92
+ # documented on this class.
93
+ self.driver.verify_connectivity()
94
+ logger.info("Neo4j driver initialized successfully")
95
+ except Exception as e:
96
+ logger.warning(f"Failed to initialize Neo4j driver: {e}")
97
+ if self.driver is not None:
98
+ try:
99
+ self.driver.close()
100
+ except Exception:
101
+ pass
102
+ self.driver = None
103
+
104
+ def close(self) -> None:
105
+ """Close the Neo4j driver connection."""
106
+ if self.driver:
107
+ self.driver.close()
108
+
109
+ def __enter__(self) -> "GraphIndex":
110
+ return self
111
+
112
+ def __exit__(self, *exc_info) -> None:
113
+ self.close()
114
+
115
+ # ------------------------------------------------------------------
116
+ # Entity extraction (regex/heuristic-based, not LLM-based — cheap
117
+ # and deterministic, matching the production approach this was
118
+ # ported from)
119
+ # ------------------------------------------------------------------
120
+
121
+ def _normalize_entity_name(self, raw_name: str) -> str:
122
+ """Normalize extracted entity text into a stable graph key.
123
+
124
+ Ported verbatim from the production implementation this package
125
+ is based on (verified via direct source review).
126
+ """
127
+ if not raw_name:
128
+ return ""
129
+ name = re.sub(r"\s+", " ", str(raw_name)).strip(" \t\r\n.,;:!?()[]{}'\"")
130
+ if len(name) < 3:
131
+ return ""
132
+ # Keep acronyms uppercase (e.g., ALARA, IAEA), title-case normal phrases.
133
+ if name.isupper():
134
+ normalized = name
135
+ else:
136
+ normalized = " ".join(part.capitalize() for part in name.split())
137
+ if len(normalized) > 120:
138
+ normalized = normalized[:120]
139
+ return normalized
140
+
141
+ def _extract_entity_candidates_from_text(
142
+ self,
143
+ text: str,
144
+ max_entities: int = 12,
145
+ domain_terms: Optional[List[str]] = None,
146
+ ) -> List[str]:
147
+ """
148
+ Lightweight, deterministic entity extraction from free text.
149
+
150
+ Ported from the production implementation this package is based
151
+ on, with one deliberate change: the source had a hardcoded list
152
+ of domain-specific terms (radiation-safety/medical-physics
153
+ vocabulary specific to that deployment's customer). That's not
154
+ appropriate as a hardcoded default in a general-purpose
155
+ open-source package — it would leak deployment-specific business
156
+ context and be irrelevant (or actively wrong) for any other
157
+ domain. Here, `domain_terms=` is an optional parameter instead —
158
+ pass your own domain vocabulary if useful, or omit it entirely
159
+ for fully generic acronym/capitalized-phrase extraction.
160
+ """
161
+ if not text:
162
+ return []
163
+
164
+ candidates: List[str] = []
165
+ # 1) Acronyms and all-caps entities.
166
+ candidates.extend(re.findall(r"\b[A-Z]{2,}(?:-[A-Z0-9]+)?\b", text))
167
+ # 2) Capitalized word sequences (person/org/concept-like phrases).
168
+ candidates.extend(
169
+ re.findall(r"\b(?:[A-Z][a-z]{2,})(?:\s+[A-Z][a-z]{2,}){0,2}\b", text)
170
+ )
171
+ # 3) Optional caller-supplied domain terms that may appear lowercase.
172
+ if domain_terms:
173
+ text_lower = text.lower()
174
+ for term in domain_terms:
175
+ if term.lower() in text_lower:
176
+ candidates.append(term)
177
+
178
+ seen = set()
179
+ final: List[str] = []
180
+ for cand in candidates:
181
+ normalized = self._normalize_entity_name(cand)
182
+ if not normalized:
183
+ continue
184
+ key = normalized.lower()
185
+ if key in seen:
186
+ continue
187
+ seen.add(key)
188
+ final.append(normalized)
189
+ if len(final) >= max_entities:
190
+ break
191
+ return final
192
+
193
+ def extract_query_entities(
194
+ self, query: str, max_entities: int = 10, domain_terms: Optional[List[str]] = None
195
+ ) -> List[str]:
196
+ """Extract entity candidates from a user query string."""
197
+ return self._extract_entity_candidates_from_text(
198
+ query, max_entities=max_entities, domain_terms=domain_terms
199
+ )
200
+
201
+ # ------------------------------------------------------------------
202
+ # Core graph operations
203
+ # ------------------------------------------------------------------
204
+
205
+ def upsert_document(
206
+ self,
207
+ document_id: str,
208
+ title: str,
209
+ chunks: List[Dict[str, Any]],
210
+ namespace: Optional[str] = None,
211
+ max_entities: int = 80,
212
+ max_pairs: int = 150,
213
+ domain_terms: Optional[List[str]] = None,
214
+ ) -> Dict[str, Any]:
215
+ """
216
+ Index one document and its extracted entities into Neo4j.
217
+
218
+ Builds co-occurrence edges between entities found in the same
219
+ chunk, weighted by how often each entity/pair appears across
220
+ the document. Writes are idempotent (MERGE-based) — safe to
221
+ re-run on the same document.
222
+
223
+ `namespace=` is optional; omit it for a single global graph, or
224
+ pass a tenant/project identifier to keep graphs isolated.
225
+
226
+ `domain_terms=` is an optional list of lowercase-matchable terms
227
+ specific to your content (e.g. industry jargon) that the default
228
+ acronym/capitalized-phrase heuristic wouldn't catch on its own.
229
+
230
+ Returns an indexing summary so callers can log/monitor coverage.
231
+ """
232
+ summary: Dict[str, Any] = {
233
+ "success": False,
234
+ "document_id": str(document_id),
235
+ "entities_indexed": 0,
236
+ "relationships_indexed": 0,
237
+ "error": None,
238
+ }
239
+
240
+ if not self.driver:
241
+ summary["error"] = "Neo4j driver not available"
242
+ return summary
243
+
244
+ entity_counter: Counter = Counter()
245
+ pair_counter: Counter = Counter()
246
+
247
+ for chunk in chunks or []:
248
+ text = (chunk.get("text") or "").strip()
249
+ if not text:
250
+ continue
251
+
252
+ entities = self._extract_entity_candidates_from_text(
253
+ text, max_entities=12, domain_terms=domain_terms
254
+ )
255
+ if not entities:
256
+ continue
257
+
258
+ unique_entities = []
259
+ seen_local = set()
260
+ for ent in entities:
261
+ key = ent.lower()
262
+ if key in seen_local:
263
+ continue
264
+ seen_local.add(key)
265
+ unique_entities.append(ent)
266
+ entity_counter[ent] += 1
267
+
268
+ for i in range(len(unique_entities)):
269
+ for j in range(i + 1, len(unique_entities)):
270
+ a, b = unique_entities[i], unique_entities[j]
271
+ if a == b:
272
+ continue
273
+ pair = tuple(sorted((a, b), key=lambda s: s.lower()))
274
+ pair_counter[pair] += 1
275
+
276
+ top_entities = entity_counter.most_common(max_entities)
277
+ top_pairs = pair_counter.most_common(max_pairs)
278
+ ns = namespace or ""
279
+
280
+ try:
281
+ with self.driver.session() as session:
282
+ session.run(
283
+ """
284
+ MERGE (d:Document {id: $document_id, namespace: $namespace})
285
+ ON CREATE SET d.created_at = datetime()
286
+ SET d.title = $title
287
+ """,
288
+ document_id=str(document_id),
289
+ namespace=ns,
290
+ title=title or "",
291
+ )
292
+
293
+ for entity_name, weight in top_entities:
294
+ session.run(
295
+ """
296
+ MATCH (d:Document {id: $document_id, namespace: $namespace})
297
+ MERGE (e:Entity {name: $name_lower, namespace: $namespace})
298
+ ON CREATE SET e.display_name = $name
299
+ MERGE (d)-[r:CONTAINS]->(e)
300
+ ON CREATE SET r.weight = $weight
301
+ ON MATCH SET r.weight = coalesce(r.weight, 0) + $weight
302
+ """,
303
+ document_id=str(document_id),
304
+ namespace=ns,
305
+ name_lower=entity_name.lower(),
306
+ name=entity_name,
307
+ weight=float(weight),
308
+ )
309
+ summary["entities_indexed"] += 1
310
+
311
+ for (a, b), weight in top_pairs:
312
+ session.run(
313
+ """
314
+ MATCH (ea:Entity {name: $a, namespace: $namespace})
315
+ MATCH (eb:Entity {name: $b, namespace: $namespace})
316
+ MERGE (ea)-[r:CO_OCCURS_WITH]-(eb)
317
+ ON CREATE SET r.weight = $weight
318
+ ON MATCH SET r.weight = coalesce(r.weight, 0) + $weight
319
+ """,
320
+ a=a.lower(),
321
+ b=b.lower(),
322
+ namespace=ns,
323
+ weight=float(weight),
324
+ )
325
+ summary["relationships_indexed"] += 1
326
+
327
+ summary["success"] = True
328
+ return summary
329
+
330
+ except Exception as e:
331
+ logger.error(f"Graph upsert error: {e}", exc_info=True)
332
+ summary["error"] = str(e)
333
+ return summary
334
+
335
+ def find_documents_by_entities(
336
+ self,
337
+ entity_names: List[str],
338
+ namespace: Optional[str] = None,
339
+ limit: int = 25,
340
+ ) -> List[Dict[str, Any]]:
341
+ """Retrieve documents connected to the given entity names, ranked
342
+ by matched-entity count and summed edge weight."""
343
+ if not self.driver:
344
+ return []
345
+
346
+ normalized = self._normalize_entity_list(entity_names)
347
+ if not normalized:
348
+ return []
349
+
350
+ ns = namespace or ""
351
+
352
+ try:
353
+ with self.driver.session() as session:
354
+ result = session.run(
355
+ """
356
+ MATCH (e:Entity)<-[r:CONTAINS]-(d:Document)
357
+ WHERE e.namespace = $namespace
358
+ AND d.namespace = $namespace
359
+ AND e.name IN $entity_names
360
+ RETURN d.id AS document_id,
361
+ coalesce(d.title, '') AS document_name,
362
+ count(DISTINCT e) AS matched_entities,
363
+ sum(coalesce(r.weight, 1.0)) AS graph_score,
364
+ collect(DISTINCT e.display_name)[0..10] AS matched_entity_names
365
+ ORDER BY matched_entities DESC, graph_score DESC
366
+ LIMIT $limit
367
+ """,
368
+ namespace=ns,
369
+ entity_names=normalized,
370
+ limit=max(1, int(limit)),
371
+ )
372
+
373
+ docs = []
374
+ for row in result:
375
+ docs.append({
376
+ "document_id": str(row["document_id"]) if row["document_id"] is not None else "",
377
+ "document_name": row["document_name"] or "",
378
+ "matched_entities": int(row["matched_entities"] or 0),
379
+ "graph_score": float(row["graph_score"] or 0.0),
380
+ "matched_entity_names": list(row["matched_entity_names"] or []),
381
+ })
382
+ return docs
383
+
384
+ except Exception as e:
385
+ logger.error(f"Graph document lookup error: {e}", exc_info=True)
386
+ return []
387
+
388
+ def search_related_entities(
389
+ self,
390
+ entity_names: List[str],
391
+ namespace: Optional[str] = None,
392
+ max_depth: int = 2,
393
+ limit: int = 10,
394
+ ) -> List[Dict[str, Any]]:
395
+ """
396
+ Find entities related to the given entity names via graph
397
+ traversal, up to `max_depth` hops.
398
+
399
+ `max_depth` is validated as a bounded integer (1-10) before use
400
+ — Neo4j's Cypher can't parameterize variable-length path bounds
401
+ normally, so this value gets string-formatted into the query.
402
+ Unvalidated input here would be a real injection vector.
403
+ """
404
+ # Validate first, regardless of driver state - invalid input is
405
+ # invalid input whether or not Neo4j happens to be reachable.
406
+ try:
407
+ max_depth = int(max_depth)
408
+ except (TypeError, ValueError):
409
+ raise ValueError(f"max_depth must be an integer, got {max_depth!r}")
410
+ if not (1 <= max_depth <= MAX_ALLOWED_DEPTH):
411
+ raise ValueError(
412
+ f"max_depth must be between 1 and {MAX_ALLOWED_DEPTH}, got {max_depth}"
413
+ )
414
+
415
+ if not self.driver:
416
+ logger.warning("Neo4j driver not available")
417
+ return []
418
+
419
+ normalized = self._normalize_entity_list(entity_names)
420
+ if not normalized:
421
+ return []
422
+
423
+ ns = namespace or ""
424
+
425
+ try:
426
+ with self.driver.session() as session:
427
+ # max_depth is validated above (bounded int) before
428
+ # being formatted in — Cypher variable-length path
429
+ # bounds can't be passed as query parameters.
430
+ query = """
431
+ MATCH (start:Entity)
432
+ WHERE start.name IN $entity_names
433
+ AND start.namespace = $namespace
434
+ MATCH path = (start)-[*1..{max_depth}]-(related:Entity)
435
+ WHERE related.namespace = $namespace
436
+ RETURN DISTINCT related.name AS entity_id,
437
+ related.display_name AS entity_name,
438
+ type(relationships(path)[0]) AS relationship,
439
+ length(path) AS depth
440
+ ORDER BY depth ASC
441
+ LIMIT $limit
442
+ """.format(max_depth=max_depth)
443
+
444
+ result = session.run(
445
+ query,
446
+ entity_names=normalized,
447
+ namespace=ns,
448
+ limit=max(1, int(limit)),
449
+ )
450
+
451
+ related = []
452
+ for row in result:
453
+ related.append({
454
+ "entity_id": row["entity_id"],
455
+ "entity_name": row["entity_name"],
456
+ "relationship": row["relationship"],
457
+ "depth": int(row["depth"]),
458
+ })
459
+ return related
460
+
461
+ except ValueError:
462
+ raise
463
+ except Exception as e:
464
+ logger.error(f"Related-entity search error: {e}", exc_info=True)
465
+ return []
466
+
467
+ def document_entities(
468
+ self,
469
+ document_id: str,
470
+ namespace: Optional[str] = None,
471
+ ) -> List[Dict[str, Any]]:
472
+ """Find all entities linked to a specific document."""
473
+ if not self.driver:
474
+ logger.warning("Neo4j driver not available")
475
+ return []
476
+
477
+ ns = namespace or ""
478
+
479
+ try:
480
+ with self.driver.session() as session:
481
+ result = session.run(
482
+ """
483
+ MATCH (d:Document {id: $document_id, namespace: $namespace})
484
+ MATCH (d)-[:CONTAINS]->(e:Entity)
485
+ RETURN e.name AS entity_id,
486
+ e.display_name AS entity_name
487
+ LIMIT 50
488
+ """,
489
+ document_id=str(document_id),
490
+ namespace=ns,
491
+ )
492
+
493
+ entities = []
494
+ for row in result:
495
+ entities.append({
496
+ "entity_id": row["entity_id"],
497
+ "entity_name": row["entity_name"],
498
+ })
499
+ return entities
500
+
501
+ except Exception as e:
502
+ logger.error(f"Document entity lookup error: {e}", exc_info=True)
503
+ return []
504
+
505
+ def health_check(self) -> bool:
506
+ """Return True if the Neo4j driver is connected and responsive."""
507
+ if not self.driver:
508
+ return False
509
+ try:
510
+ with self.driver.session() as session:
511
+ session.run("RETURN 1")
512
+ return True
513
+ except Exception as e:
514
+ logger.warning(f"Graph health check failed: {e}")
515
+ return False
516
+
517
+ # ------------------------------------------------------------------
518
+ # Internal helpers
519
+ # ------------------------------------------------------------------
520
+
521
+ def _normalize_entity_list(self, entity_names: List[str]) -> List[str]:
522
+ normalized = []
523
+ seen = set()
524
+ for name in entity_names or []:
525
+ norm = self._normalize_entity_name(name)
526
+ if not norm:
527
+ continue
528
+ key = norm.lower()
529
+ if key in seen:
530
+ continue
531
+ seen.add(key)
532
+ normalized.append(key)
533
+ return normalized
534
+
535
+
536
+ __all__ = ["GraphConfig", "GraphIndex", "__version__"]
@@ -0,0 +1,251 @@
1
+ """
2
+ Tests for ragleap_graph.
3
+
4
+ Three tiers, mirroring ragleap-rag's own test discipline:
5
+ 1. Pure logic tests — no Neo4j needed (normalization, extraction, validation)
6
+ 2. Driver-unavailable safety tests — confirm graceful degradation
7
+ 3. Live integration test — real Neo4j round-trip, skipped automatically
8
+ unless NEO4J_URI/NEO4J_USER/NEO4J_PASSWORD are set in the environment.
9
+ Cleans up its own test data via namespace-scoped delete, same pattern
10
+ manually verified during development.
11
+ """
12
+ import os
13
+
14
+ import pytest
15
+
16
+ from ragleap_graph import GraphConfig, GraphIndex, MAX_ALLOWED_DEPTH
17
+
18
+
19
+ # ---------------------------------------------------------------------------
20
+ # GraphConfig
21
+ # ---------------------------------------------------------------------------
22
+
23
+ def test_graph_config_defaults():
24
+ config = GraphConfig()
25
+ assert config.uri == "bolt://localhost:7687"
26
+ assert config.user == "neo4j"
27
+ assert config.password == ""
28
+
29
+
30
+ def test_graph_config_custom_values():
31
+ config = GraphConfig(uri="bolt://example.com:7687", user="admin", password="secret")
32
+ assert config.uri == "bolt://example.com:7687"
33
+ assert config.user == "admin"
34
+ assert config.password == "secret"
35
+
36
+
37
+ # ---------------------------------------------------------------------------
38
+ # Entity normalization (pure logic, no driver needed)
39
+ # ---------------------------------------------------------------------------
40
+
41
+ @pytest.fixture
42
+ def graph_no_driver():
43
+ """A GraphIndex whose driver never connects (bad URI) — safe for
44
+ testing pure logic methods without touching a real database."""
45
+ config = GraphConfig(uri="bolt://localhost:1", user="x", password="x")
46
+ g = GraphIndex(config=config)
47
+ assert g.driver is None
48
+ return g
49
+
50
+
51
+ def test_normalize_entity_name_empty(graph_no_driver):
52
+ assert graph_no_driver._normalize_entity_name("") == ""
53
+ assert graph_no_driver._normalize_entity_name(None) == ""
54
+
55
+
56
+ def test_normalize_entity_name_too_short(graph_no_driver):
57
+ assert graph_no_driver._normalize_entity_name("ab") == ""
58
+
59
+
60
+ def test_normalize_entity_name_collapses_whitespace(graph_no_driver):
61
+ result = graph_no_driver._normalize_entity_name("Acme Corp\n\tInc")
62
+ assert result == "Acme Corp Inc"
63
+
64
+
65
+ def test_normalize_entity_name_strips_punctuation(graph_no_driver):
66
+ result = graph_no_driver._normalize_entity_name(" (Acme Corp). ")
67
+ assert result == "Acme Corp"
68
+
69
+
70
+ def test_normalize_entity_name_keeps_acronyms_uppercase(graph_no_driver):
71
+ assert graph_no_driver._normalize_entity_name("ALARA") == "ALARA"
72
+ assert graph_no_driver._normalize_entity_name("IAEA") == "IAEA"
73
+
74
+
75
+ def test_normalize_entity_name_title_cases_normal_phrases(graph_no_driver):
76
+ assert graph_no_driver._normalize_entity_name("acme corp") == "Acme Corp"
77
+ assert graph_no_driver._normalize_entity_name("ACME corp") == "Acme Corp"
78
+
79
+
80
+ def test_normalize_entity_name_truncates_long_names(graph_no_driver):
81
+ long_name = "A" * 200
82
+ result = graph_no_driver._normalize_entity_name(long_name)
83
+ assert len(result) == 120
84
+
85
+
86
+ # ---------------------------------------------------------------------------
87
+ # Entity extraction (pure logic, no driver needed)
88
+ # ---------------------------------------------------------------------------
89
+
90
+ def test_extract_entities_empty_text(graph_no_driver):
91
+ assert graph_no_driver._extract_entity_candidates_from_text("") == []
92
+
93
+
94
+ def test_extract_entities_finds_acronyms(graph_no_driver):
95
+ result = graph_no_driver._extract_entity_candidates_from_text(
96
+ "The IAEA published new guidelines."
97
+ )
98
+ assert "IAEA" in result
99
+
100
+
101
+ def test_extract_entities_finds_capitalized_phrases(graph_no_driver):
102
+ result = graph_no_driver._extract_entity_candidates_from_text(
103
+ "Acme Corp signed a deal with Globex Industries."
104
+ )
105
+ assert "Acme Corp" in result
106
+ assert "Globex Industries" in result
107
+
108
+
109
+ def test_extract_entities_respects_max_entities(graph_no_driver):
110
+ text = " ".join(f"Company{i} Corp" for i in range(20))
111
+ result = graph_no_driver._extract_entity_candidates_from_text(text, max_entities=5)
112
+ assert len(result) <= 5
113
+
114
+
115
+ def test_extract_entities_deduplicates(graph_no_driver):
116
+ result = graph_no_driver._extract_entity_candidates_from_text(
117
+ "Acme Corp is great. Acme Corp is really great."
118
+ )
119
+ assert result.count("Acme Corp") == 1
120
+
121
+
122
+ def test_extract_entities_domain_terms_optional_and_off_by_default(graph_no_driver):
123
+ """A term that only matches via domain_terms shouldn't appear unless
124
+ domain_terms is explicitly passed — no hardcoded vocabulary leaks in."""
125
+ text = "we discussed radiation safety protocols today"
126
+ without = graph_no_driver._extract_entity_candidates_from_text(text)
127
+ assert not any("radiation" in e.lower() for e in without)
128
+
129
+ with_terms = graph_no_driver._extract_entity_candidates_from_text(
130
+ text, domain_terms=["radiation safety"]
131
+ )
132
+ assert any("radiation safety" in e.lower() for e in with_terms)
133
+
134
+
135
+ def test_extract_query_entities_delegates_correctly(graph_no_driver):
136
+ result = graph_no_driver.extract_query_entities("What does Acme Corp do?")
137
+ assert "Acme Corp" in result
138
+
139
+
140
+ # ---------------------------------------------------------------------------
141
+ # max_depth validation (the security fix — real Cypher-injection guard)
142
+ # ---------------------------------------------------------------------------
143
+
144
+ def test_max_depth_rejects_non_integer(graph_no_driver):
145
+ with pytest.raises(ValueError):
146
+ graph_no_driver.search_related_entities(["x"], max_depth="1; DROP")
147
+
148
+
149
+ def test_max_depth_rejects_too_large(graph_no_driver):
150
+ with pytest.raises(ValueError):
151
+ graph_no_driver.search_related_entities(["x"], max_depth=MAX_ALLOWED_DEPTH + 1)
152
+
153
+
154
+ def test_max_depth_rejects_zero_or_negative(graph_no_driver):
155
+ with pytest.raises(ValueError):
156
+ graph_no_driver.search_related_entities(["x"], max_depth=0)
157
+ with pytest.raises(ValueError):
158
+ graph_no_driver.search_related_entities(["x"], max_depth=-1)
159
+
160
+
161
+ def test_max_depth_accepts_valid_range(graph_no_driver):
162
+ # driver is None, so this returns [] rather than raising — the point
163
+ # here is only that validation itself doesn't reject valid input.
164
+ result = graph_no_driver.search_related_entities(["x"], max_depth=3)
165
+ assert result == []
166
+
167
+
168
+ # ---------------------------------------------------------------------------
169
+ # Driver-unavailable safety (every public method must degrade gracefully)
170
+ # ---------------------------------------------------------------------------
171
+
172
+ def test_health_check_false_without_driver(graph_no_driver):
173
+ assert graph_no_driver.health_check() is False
174
+
175
+
176
+ def test_upsert_document_returns_error_without_driver(graph_no_driver):
177
+ result = graph_no_driver.upsert_document(
178
+ document_id="d1", title="t", chunks=[{"text": "Acme Corp"}]
179
+ )
180
+ assert result["success"] is False
181
+ assert "not available" in result["error"]
182
+
183
+
184
+ def test_find_documents_empty_without_driver(graph_no_driver):
185
+ assert graph_no_driver.find_documents_by_entities(["Acme Corp"]) == []
186
+
187
+
188
+ def test_document_entities_empty_without_driver(graph_no_driver):
189
+ assert graph_no_driver.document_entities("doc-1") == []
190
+
191
+
192
+ def test_context_manager_closes_cleanly(graph_no_driver):
193
+ with graph_no_driver as g:
194
+ assert g is graph_no_driver
195
+
196
+
197
+ # ---------------------------------------------------------------------------
198
+ # Live integration test — real Neo4j, real round-trip, self-cleaning.
199
+ # ---------------------------------------------------------------------------
200
+
201
+ NEO4J_URI = os.environ.get("NEO4J_URI")
202
+ NEO4J_USER = os.environ.get("NEO4J_USER")
203
+ NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD")
204
+ HAS_LIVE_NEO4J = bool(NEO4J_URI and NEO4J_USER and NEO4J_PASSWORD)
205
+
206
+ TEST_NAMESPACE = "ragleap_graph_pytest"
207
+
208
+
209
+ @pytest.mark.skipif(not HAS_LIVE_NEO4J, reason="No live Neo4j credentials in environment")
210
+ def test_live_full_roundtrip():
211
+ """
212
+ Real end-to-end test against a real Neo4j instance: connect, write,
213
+ query three ways, then clean up every node this test created —
214
+ never leaves data behind in a real database.
215
+ """
216
+ config = GraphConfig(uri=NEO4J_URI, user=NEO4J_USER, password=NEO4J_PASSWORD)
217
+ graph = GraphIndex(config=config)
218
+
219
+ try:
220
+ assert graph.health_check() is True
221
+
222
+ result = graph.upsert_document(
223
+ document_id="pytest-live-doc-1",
224
+ title="Pytest Live Test Document",
225
+ chunks=[{"text": "Acme Corp announced a partnership with Globex Industries."}],
226
+ namespace=TEST_NAMESPACE,
227
+ )
228
+ assert result["success"] is True
229
+ assert result["entities_indexed"] > 0
230
+
231
+ docs = graph.find_documents_by_entities(["Acme Corp"], namespace=TEST_NAMESPACE)
232
+ assert len(docs) == 1
233
+ assert docs[0]["document_id"] == "pytest-live-doc-1"
234
+
235
+ related = graph.search_related_entities(
236
+ ["Acme Corp"], namespace=TEST_NAMESPACE, max_depth=2
237
+ )
238
+ assert any(r["entity_name"] == "Globex Industries" for r in related)
239
+
240
+ entities = graph.document_entities("pytest-live-doc-1", namespace=TEST_NAMESPACE)
241
+ entity_names = {e["entity_name"] for e in entities}
242
+ assert "Acme Corp" in entity_names
243
+ assert "Globex Industries" in entity_names
244
+
245
+ finally:
246
+ with graph.driver.session() as session:
247
+ session.run(
248
+ "MATCH (n) WHERE n.namespace = $ns DETACH DELETE n",
249
+ ns=TEST_NAMESPACE,
250
+ )
251
+ graph.close()