simple-graph-retriever 1.0.0rc1__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.
- simple_graph_retriever-1.0.0rc1/.github/workflows/release.yml +33 -0
- simple_graph_retriever-1.0.0rc1/.gitignore +6 -0
- simple_graph_retriever-1.0.0rc1/.releaserc.json +62 -0
- simple_graph_retriever-1.0.0rc1/PKG-INFO +203 -0
- simple_graph_retriever-1.0.0rc1/README.md +184 -0
- simple_graph_retriever-1.0.0rc1/examples/index.py +26 -0
- simple_graph_retriever-1.0.0rc1/examples/retrieve.py +55 -0
- simple_graph_retriever-1.0.0rc1/pyproject.toml +31 -0
- simple_graph_retriever-1.0.0rc1/simple_graph_retriever/__init__.py +0 -0
- simple_graph_retriever-1.0.0rc1/simple_graph_retriever/client.py +78 -0
- simple_graph_retriever-1.0.0rc1/simple_graph_retriever/config.py +36 -0
- simple_graph_retriever-1.0.0rc1/simple_graph_retriever/indexer.py +367 -0
- simple_graph_retriever-1.0.0rc1/simple_graph_retriever/models.py +16 -0
- simple_graph_retriever-1.0.0rc1/simple_graph_retriever/retriever.py +244 -0
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
name: Release
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches:
|
|
6
|
+
- main
|
|
7
|
+
- release-candidate
|
|
8
|
+
|
|
9
|
+
jobs:
|
|
10
|
+
release:
|
|
11
|
+
permissions:
|
|
12
|
+
contents: write
|
|
13
|
+
issues: write
|
|
14
|
+
runs-on: ubuntu-latest
|
|
15
|
+
steps:
|
|
16
|
+
- uses: actions/checkout@v4
|
|
17
|
+
- name: Set up Python
|
|
18
|
+
uses: actions/setup-python@v2
|
|
19
|
+
with:
|
|
20
|
+
python-version: "3.x"
|
|
21
|
+
- name: Install dependencies
|
|
22
|
+
run: |
|
|
23
|
+
python -m pip install --upgrade pip
|
|
24
|
+
pip install wheel setuptools build
|
|
25
|
+
npm install -g semantic-release semantic-release-pypi @semantic-release/github @semantic-release/git semantic-release-replace-plugin
|
|
26
|
+
- name: Clean dist directory
|
|
27
|
+
run: rm -rf dist/
|
|
28
|
+
- name: Release
|
|
29
|
+
env:
|
|
30
|
+
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
|
31
|
+
PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
|
|
32
|
+
PYPI_REPOSITORY_URL: https://test.pypi.org/legacy/
|
|
33
|
+
run: npx semantic-release
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
{
|
|
2
|
+
"branches": [
|
|
3
|
+
"main",
|
|
4
|
+
{
|
|
5
|
+
"name": "release-candidate",
|
|
6
|
+
"prerelease": "rc"
|
|
7
|
+
}
|
|
8
|
+
],
|
|
9
|
+
"plugins": [
|
|
10
|
+
"@semantic-release/commit-analyzer",
|
|
11
|
+
"@semantic-release/release-notes-generator",
|
|
12
|
+
[
|
|
13
|
+
"semantic-release-pypi",
|
|
14
|
+
{
|
|
15
|
+
"repoUrl": "https://upload.pypi.org/legacy/"
|
|
16
|
+
}
|
|
17
|
+
],
|
|
18
|
+
[
|
|
19
|
+
"@semantic-release/github",
|
|
20
|
+
{
|
|
21
|
+
"assets": [
|
|
22
|
+
{
|
|
23
|
+
"path": "dist/*.whl",
|
|
24
|
+
"label": "Wheel distribution"
|
|
25
|
+
},
|
|
26
|
+
{
|
|
27
|
+
"path": "dist/*.tar.gz",
|
|
28
|
+
"label": "Source distribution"
|
|
29
|
+
}
|
|
30
|
+
]
|
|
31
|
+
}
|
|
32
|
+
],
|
|
33
|
+
[
|
|
34
|
+
"semantic-release-replace-plugin",
|
|
35
|
+
{
|
|
36
|
+
"replacements": [
|
|
37
|
+
{
|
|
38
|
+
"files": ["README.md"],
|
|
39
|
+
"from": "pip install simple_graph_retriever==.*",
|
|
40
|
+
"to": "pip install simple_graph_retriever==${nextRelease.version}",
|
|
41
|
+
"results": [
|
|
42
|
+
{
|
|
43
|
+
"file": "README.md",
|
|
44
|
+
"hasChanged": true,
|
|
45
|
+
"numMatches": 1,
|
|
46
|
+
"numReplacements": 1
|
|
47
|
+
}
|
|
48
|
+
],
|
|
49
|
+
"countMatches": true
|
|
50
|
+
}
|
|
51
|
+
]
|
|
52
|
+
}
|
|
53
|
+
],
|
|
54
|
+
[
|
|
55
|
+
"@semantic-release/git",
|
|
56
|
+
{
|
|
57
|
+
"assets": ["pyproject.toml", "README.md"],
|
|
58
|
+
"message": "chore(release): ${nextRelease.version} [skip ci]\n\n${nextRelease.notes}"
|
|
59
|
+
}
|
|
60
|
+
]
|
|
61
|
+
]
|
|
62
|
+
}
|
|
@@ -0,0 +1,203 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: simple-graph-retriever
|
|
3
|
+
Version: 1.0.0rc1
|
|
4
|
+
Summary: A simple graph-based retriever using Neo4j and Qdrant
|
|
5
|
+
Author-email: Victor de La Salmonière <victor@lettria.com>
|
|
6
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
7
|
+
Classifier: Operating System :: OS Independent
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Requires-Python: >=3.8
|
|
10
|
+
Requires-Dist: igraph
|
|
11
|
+
Requires-Dist: leidenalg
|
|
12
|
+
Requires-Dist: neo4j
|
|
13
|
+
Requires-Dist: pydantic
|
|
14
|
+
Requires-Dist: pydantic-settings
|
|
15
|
+
Requires-Dist: qdrant-client
|
|
16
|
+
Requires-Dist: requests
|
|
17
|
+
Requires-Dist: tqdm
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
|
|
20
|
+
# Simple Graph Retriever
|
|
21
|
+
|
|
22
|
+
A Python SDK for indexing a Neo4j graph into a vector database (Qdrant) and performing retrieval-augmented generation (RAG) tasks against it.
|
|
23
|
+
|
|
24
|
+
This SDK handles the heavy lifting of graph processing, including:
|
|
25
|
+
|
|
26
|
+
- **Community Detection**: Uses the Leiden algorithm to identify communities within your graph structure.
|
|
27
|
+
- **Chunking**: Breaks down nodes and their local context into "chunks" suitable for embedding.
|
|
28
|
+
- **Vector Indexing**: Embeds and indexes both community summaries and individual chunks into Qdrant.
|
|
29
|
+
- **Graph Retrieval**: Retrieves relevant subgraphs based on a natural language query, with fine-grained control over the expansion process.
|
|
30
|
+
|
|
31
|
+
---
|
|
32
|
+
|
|
33
|
+
## Table of Contents
|
|
34
|
+
|
|
35
|
+
- [Simple Graph Retriever](#simple-graph-retriever)
|
|
36
|
+
- [Table of Contents](#table-of-contents)
|
|
37
|
+
- [1. Installation](#1-installation)
|
|
38
|
+
- [2. Configuration](#2-configuration)
|
|
39
|
+
- [3. Usage](#3-usage)
|
|
40
|
+
- [Initializing the Client](#initializing-the-client)
|
|
41
|
+
- [Indexing the Graph](#indexing-the-graph)
|
|
42
|
+
- [Retrieving a Subgraph](#retrieving-a-subgraph)
|
|
43
|
+
- [Basic Retrieval](#basic-retrieval)
|
|
44
|
+
- [Advanced Retrieval with Graph Expansion Control](#advanced-retrieval-with-graph-expansion-control)
|
|
45
|
+
- [4. Project Structure](#4-project-structure)
|
|
46
|
+
- [5. Contributing](#5-contributing)
|
|
47
|
+
- [6. License](#6-license)
|
|
48
|
+
|
|
49
|
+
---
|
|
50
|
+
|
|
51
|
+
## 1. Installation
|
|
52
|
+
|
|
53
|
+
To install the SDK, clone this repository and use pip to install it in editable mode, which is recommended for development.
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
pip install simple_graph_retriever==0.0.0
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
## 2. Configuration
|
|
60
|
+
|
|
61
|
+
The SDK is configured via environment variables, which are loaded from a `.env` file in your project's root directory.
|
|
62
|
+
|
|
63
|
+
| Variable | Default Value | Description |
|
|
64
|
+
| ------------------------------- | ----------------------- | ------------------------------------------------------------------------------------------------- |
|
|
65
|
+
| `NEO4J_URI` | `bolt://localhost:7687` | The URI for your Neo4j database. |
|
|
66
|
+
| `NEO4J_USER` | `neo4j` | The username for your Neo4j database. |
|
|
67
|
+
| `NEO4J_PASSWORD` | `password` | The password for your Neo4j database. |
|
|
68
|
+
| `QDRANT_URL` | `http://localhost:6333` | The URL for your Qdrant vector database instance. |
|
|
69
|
+
| `QDRANT_API_KEY` | `None` | Optional: The API key for authenticating with your Qdrant instance. |
|
|
70
|
+
| `QDRANT_CHUNKS_COLLECTION` | `chunks` | The name of the collection for storing graph chunks. |
|
|
71
|
+
| `QDRANT_COMMUNITIES_COLLECTION` | `communities` | The name of the collection for storing community summaries. |
|
|
72
|
+
| `EMBEDDER_URL` | `http://localhost:8080` | The URL of the text embedding service. It must accept a POST request with `{"inputs": ["text"]}`. |
|
|
73
|
+
| `VECTOR_SIZE` | `384` | The dimension of the vectors produced by your embedding model. |
|
|
74
|
+
| `LOGLEVEL` | `INFO` | The logging level for the SDK (`DEBUG`, `INFO`, `WARNING`, `ERROR`). |
|
|
75
|
+
|
|
76
|
+
**Example `.env` file:**
|
|
77
|
+
|
|
78
|
+
```dotenv
|
|
79
|
+
# Neo4j Database Credentials
|
|
80
|
+
NEO4J_URI="bolt://localhost:7687"
|
|
81
|
+
NEO4J_USER="neo4j"
|
|
82
|
+
NEO4J_PASSWORD="your_secure_password"
|
|
83
|
+
|
|
84
|
+
# Qdrant Vector Database URL
|
|
85
|
+
QDRANT_URL="http://localhost:6333"
|
|
86
|
+
QDRANT_API_KEY="your_qdrant_api_key" # Optional
|
|
87
|
+
QDRANT_CHUNKS_COLLECTION="chunks"
|
|
88
|
+
QDRANT_COMMUNITIES_COLLECTION="communities"
|
|
89
|
+
|
|
90
|
+
# Text Embedding Model URL
|
|
91
|
+
EMBEDDER_URL="http://localhost:8080/embed"
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
## 3. Usage
|
|
95
|
+
|
|
96
|
+
### Initializing the Client
|
|
97
|
+
|
|
98
|
+
The main entry point is the `GraphRetrievalClient`. It automatically loads settings from your `.env` file.
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
from graph_retrieval_sdk.client import GraphRetrievalClient
|
|
102
|
+
|
|
103
|
+
# Initialize the client
|
|
104
|
+
client = GraphRetrievalClient()
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
### Indexing the Graph
|
|
108
|
+
|
|
109
|
+
The `index()` method runs the full, idempotent pipeline to populate Qdrant with your graph data.
|
|
110
|
+
|
|
111
|
+
```python
|
|
112
|
+
# This runs the full pipeline:
|
|
113
|
+
# 1. Community detection
|
|
114
|
+
# 2. Chunk creation
|
|
115
|
+
# 3. Chunk and Community indexing
|
|
116
|
+
client.index()
|
|
117
|
+
|
|
118
|
+
print("Indexing complete!")
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
To completely reset your vector index, use the `clear_index()` method. **Warning**: This is a destructive operation that deletes and recreates the Qdrant collections.
|
|
122
|
+
|
|
123
|
+
```python
|
|
124
|
+
client.clear_index()
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
### Retrieving a Subgraph
|
|
128
|
+
|
|
129
|
+
Use the `retrieve_graph()` method to query your indexed graph. The `RetrievalConfig` model allows for fine-grained control over the process.
|
|
130
|
+
|
|
131
|
+
#### Basic Retrieval
|
|
132
|
+
|
|
133
|
+
```python
|
|
134
|
+
from graph_retrieval_sdk.models import RetrievalConfig
|
|
135
|
+
import json
|
|
136
|
+
|
|
137
|
+
query = "Who was the emperor of Rome?"
|
|
138
|
+
config = RetrievalConfig()
|
|
139
|
+
|
|
140
|
+
subgraph = client.retriever.retrieve_graph(query, config)
|
|
141
|
+
|
|
142
|
+
if subgraph:
|
|
143
|
+
print(f"Retrieved {len(subgraph[0]['nodes'])} nodes and {len(subgraph[0]['relationships'])} relationships.")
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
#### Advanced Retrieval with Graph Expansion Control
|
|
147
|
+
|
|
148
|
+
Customize the `RetrievalConfig` to tune the retrieval and expansion behavior.
|
|
149
|
+
|
|
150
|
+
```python
|
|
151
|
+
from graph_retrieval_sdk.models import RetrievalConfig
|
|
152
|
+
|
|
153
|
+
query = "Tell me about the Roman emperors and their families, but exclude servants."
|
|
154
|
+
|
|
155
|
+
advanced_config = RetrievalConfig(
|
|
156
|
+
# --- Retrieval Settings ---
|
|
157
|
+
max_communities=5,
|
|
158
|
+
max_chunks=15,
|
|
159
|
+
|
|
160
|
+
# --- Graph Expansion Settings ---
|
|
161
|
+
max_hops=2,
|
|
162
|
+
community_expansion_limit=5,
|
|
163
|
+
|
|
164
|
+
# --- Relationship Filtering ---
|
|
165
|
+
allowed_rel_types=["HAS_SON", "MARRIED_TO", "SUCCESSOR_OF"],
|
|
166
|
+
denied_rel_types=["HAS_SERVANT"]
|
|
167
|
+
)
|
|
168
|
+
|
|
169
|
+
subgraph = client.retriever.retrieve_graph(query, advanced_config)
|
|
170
|
+
|
|
171
|
+
if subgraph:
|
|
172
|
+
print(f"Retrieved {len(subgraph[0]['nodes'])} nodes with advanced configuration.")
|
|
173
|
+
|
|
174
|
+
# Close the client connection when done
|
|
175
|
+
client.close()
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
## 4. Project Structure
|
|
179
|
+
|
|
180
|
+
The project is organized as a standard Python package:
|
|
181
|
+
|
|
182
|
+
```
|
|
183
|
+
├── graph_retrieval_sdk/
|
|
184
|
+
│ ├── __init__.py
|
|
185
|
+
│ ├── client.py # Main client class
|
|
186
|
+
│ ├── config.py # Configuration and logging setup
|
|
187
|
+
│ ├── indexer.py # Logic for indexing the graph
|
|
188
|
+
│ ├── retriever.py # Logic for retrieving subgraphs
|
|
189
|
+
│ └── models.py # Pydantic models (e.g., RetrievalConfig)
|
|
190
|
+
├── examples/
|
|
191
|
+
│ ├── retrieve_graph.py
|
|
192
|
+
│ └── run_indexing.py
|
|
193
|
+
├── pyproject.toml # Project metadata and dependencies
|
|
194
|
+
└── README.md # This file
|
|
195
|
+
```
|
|
196
|
+
|
|
197
|
+
## 5. Contributing
|
|
198
|
+
|
|
199
|
+
Contributions are welcome! Please feel free to submit a pull request or open an issue.
|
|
200
|
+
|
|
201
|
+
## 6. License
|
|
202
|
+
|
|
203
|
+
This project is licensed under the MIT License.
|
|
@@ -0,0 +1,184 @@
|
|
|
1
|
+
# Simple Graph Retriever
|
|
2
|
+
|
|
3
|
+
A Python SDK for indexing a Neo4j graph into a vector database (Qdrant) and performing retrieval-augmented generation (RAG) tasks against it.
|
|
4
|
+
|
|
5
|
+
This SDK handles the heavy lifting of graph processing, including:
|
|
6
|
+
|
|
7
|
+
- **Community Detection**: Uses the Leiden algorithm to identify communities within your graph structure.
|
|
8
|
+
- **Chunking**: Breaks down nodes and their local context into "chunks" suitable for embedding.
|
|
9
|
+
- **Vector Indexing**: Embeds and indexes both community summaries and individual chunks into Qdrant.
|
|
10
|
+
- **Graph Retrieval**: Retrieves relevant subgraphs based on a natural language query, with fine-grained control over the expansion process.
|
|
11
|
+
|
|
12
|
+
---
|
|
13
|
+
|
|
14
|
+
## Table of Contents
|
|
15
|
+
|
|
16
|
+
- [Simple Graph Retriever](#simple-graph-retriever)
|
|
17
|
+
- [Table of Contents](#table-of-contents)
|
|
18
|
+
- [1. Installation](#1-installation)
|
|
19
|
+
- [2. Configuration](#2-configuration)
|
|
20
|
+
- [3. Usage](#3-usage)
|
|
21
|
+
- [Initializing the Client](#initializing-the-client)
|
|
22
|
+
- [Indexing the Graph](#indexing-the-graph)
|
|
23
|
+
- [Retrieving a Subgraph](#retrieving-a-subgraph)
|
|
24
|
+
- [Basic Retrieval](#basic-retrieval)
|
|
25
|
+
- [Advanced Retrieval with Graph Expansion Control](#advanced-retrieval-with-graph-expansion-control)
|
|
26
|
+
- [4. Project Structure](#4-project-structure)
|
|
27
|
+
- [5. Contributing](#5-contributing)
|
|
28
|
+
- [6. License](#6-license)
|
|
29
|
+
|
|
30
|
+
---
|
|
31
|
+
|
|
32
|
+
## 1. Installation
|
|
33
|
+
|
|
34
|
+
To install the SDK, clone this repository and use pip to install it in editable mode, which is recommended for development.
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install simple_graph_retriever==0.0.0
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
## 2. Configuration
|
|
41
|
+
|
|
42
|
+
The SDK is configured via environment variables, which are loaded from a `.env` file in your project's root directory.
|
|
43
|
+
|
|
44
|
+
| Variable | Default Value | Description |
|
|
45
|
+
| ------------------------------- | ----------------------- | ------------------------------------------------------------------------------------------------- |
|
|
46
|
+
| `NEO4J_URI` | `bolt://localhost:7687` | The URI for your Neo4j database. |
|
|
47
|
+
| `NEO4J_USER` | `neo4j` | The username for your Neo4j database. |
|
|
48
|
+
| `NEO4J_PASSWORD` | `password` | The password for your Neo4j database. |
|
|
49
|
+
| `QDRANT_URL` | `http://localhost:6333` | The URL for your Qdrant vector database instance. |
|
|
50
|
+
| `QDRANT_API_KEY` | `None` | Optional: The API key for authenticating with your Qdrant instance. |
|
|
51
|
+
| `QDRANT_CHUNKS_COLLECTION` | `chunks` | The name of the collection for storing graph chunks. |
|
|
52
|
+
| `QDRANT_COMMUNITIES_COLLECTION` | `communities` | The name of the collection for storing community summaries. |
|
|
53
|
+
| `EMBEDDER_URL` | `http://localhost:8080` | The URL of the text embedding service. It must accept a POST request with `{"inputs": ["text"]}`. |
|
|
54
|
+
| `VECTOR_SIZE` | `384` | The dimension of the vectors produced by your embedding model. |
|
|
55
|
+
| `LOGLEVEL` | `INFO` | The logging level for the SDK (`DEBUG`, `INFO`, `WARNING`, `ERROR`). |
|
|
56
|
+
|
|
57
|
+
**Example `.env` file:**
|
|
58
|
+
|
|
59
|
+
```dotenv
|
|
60
|
+
# Neo4j Database Credentials
|
|
61
|
+
NEO4J_URI="bolt://localhost:7687"
|
|
62
|
+
NEO4J_USER="neo4j"
|
|
63
|
+
NEO4J_PASSWORD="your_secure_password"
|
|
64
|
+
|
|
65
|
+
# Qdrant Vector Database URL
|
|
66
|
+
QDRANT_URL="http://localhost:6333"
|
|
67
|
+
QDRANT_API_KEY="your_qdrant_api_key" # Optional
|
|
68
|
+
QDRANT_CHUNKS_COLLECTION="chunks"
|
|
69
|
+
QDRANT_COMMUNITIES_COLLECTION="communities"
|
|
70
|
+
|
|
71
|
+
# Text Embedding Model URL
|
|
72
|
+
EMBEDDER_URL="http://localhost:8080/embed"
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
## 3. Usage
|
|
76
|
+
|
|
77
|
+
### Initializing the Client
|
|
78
|
+
|
|
79
|
+
The main entry point is the `GraphRetrievalClient`. It automatically loads settings from your `.env` file.
|
|
80
|
+
|
|
81
|
+
```python
|
|
82
|
+
from graph_retrieval_sdk.client import GraphRetrievalClient
|
|
83
|
+
|
|
84
|
+
# Initialize the client
|
|
85
|
+
client = GraphRetrievalClient()
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
### Indexing the Graph
|
|
89
|
+
|
|
90
|
+
The `index()` method runs the full, idempotent pipeline to populate Qdrant with your graph data.
|
|
91
|
+
|
|
92
|
+
```python
|
|
93
|
+
# This runs the full pipeline:
|
|
94
|
+
# 1. Community detection
|
|
95
|
+
# 2. Chunk creation
|
|
96
|
+
# 3. Chunk and Community indexing
|
|
97
|
+
client.index()
|
|
98
|
+
|
|
99
|
+
print("Indexing complete!")
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
To completely reset your vector index, use the `clear_index()` method. **Warning**: This is a destructive operation that deletes and recreates the Qdrant collections.
|
|
103
|
+
|
|
104
|
+
```python
|
|
105
|
+
client.clear_index()
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
### Retrieving a Subgraph
|
|
109
|
+
|
|
110
|
+
Use the `retrieve_graph()` method to query your indexed graph. The `RetrievalConfig` model allows for fine-grained control over the process.
|
|
111
|
+
|
|
112
|
+
#### Basic Retrieval
|
|
113
|
+
|
|
114
|
+
```python
|
|
115
|
+
from graph_retrieval_sdk.models import RetrievalConfig
|
|
116
|
+
import json
|
|
117
|
+
|
|
118
|
+
query = "Who was the emperor of Rome?"
|
|
119
|
+
config = RetrievalConfig()
|
|
120
|
+
|
|
121
|
+
subgraph = client.retriever.retrieve_graph(query, config)
|
|
122
|
+
|
|
123
|
+
if subgraph:
|
|
124
|
+
print(f"Retrieved {len(subgraph[0]['nodes'])} nodes and {len(subgraph[0]['relationships'])} relationships.")
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
#### Advanced Retrieval with Graph Expansion Control
|
|
128
|
+
|
|
129
|
+
Customize the `RetrievalConfig` to tune the retrieval and expansion behavior.
|
|
130
|
+
|
|
131
|
+
```python
|
|
132
|
+
from graph_retrieval_sdk.models import RetrievalConfig
|
|
133
|
+
|
|
134
|
+
query = "Tell me about the Roman emperors and their families, but exclude servants."
|
|
135
|
+
|
|
136
|
+
advanced_config = RetrievalConfig(
|
|
137
|
+
# --- Retrieval Settings ---
|
|
138
|
+
max_communities=5,
|
|
139
|
+
max_chunks=15,
|
|
140
|
+
|
|
141
|
+
# --- Graph Expansion Settings ---
|
|
142
|
+
max_hops=2,
|
|
143
|
+
community_expansion_limit=5,
|
|
144
|
+
|
|
145
|
+
# --- Relationship Filtering ---
|
|
146
|
+
allowed_rel_types=["HAS_SON", "MARRIED_TO", "SUCCESSOR_OF"],
|
|
147
|
+
denied_rel_types=["HAS_SERVANT"]
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
subgraph = client.retriever.retrieve_graph(query, advanced_config)
|
|
151
|
+
|
|
152
|
+
if subgraph:
|
|
153
|
+
print(f"Retrieved {len(subgraph[0]['nodes'])} nodes with advanced configuration.")
|
|
154
|
+
|
|
155
|
+
# Close the client connection when done
|
|
156
|
+
client.close()
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
## 4. Project Structure
|
|
160
|
+
|
|
161
|
+
The project is organized as a standard Python package:
|
|
162
|
+
|
|
163
|
+
```
|
|
164
|
+
├── graph_retrieval_sdk/
|
|
165
|
+
│ ├── __init__.py
|
|
166
|
+
│ ├── client.py # Main client class
|
|
167
|
+
│ ├── config.py # Configuration and logging setup
|
|
168
|
+
│ ├── indexer.py # Logic for indexing the graph
|
|
169
|
+
│ ├── retriever.py # Logic for retrieving subgraphs
|
|
170
|
+
│ └── models.py # Pydantic models (e.g., RetrievalConfig)
|
|
171
|
+
├── examples/
|
|
172
|
+
│ ├── retrieve_graph.py
|
|
173
|
+
│ └── run_indexing.py
|
|
174
|
+
├── pyproject.toml # Project metadata and dependencies
|
|
175
|
+
└── README.md # This file
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
## 5. Contributing
|
|
179
|
+
|
|
180
|
+
Contributions are welcome! Please feel free to submit a pull request or open an issue.
|
|
181
|
+
|
|
182
|
+
## 6. License
|
|
183
|
+
|
|
184
|
+
This project is licensed under the MIT License.
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
import sys
|
|
2
|
+
import os
|
|
3
|
+
|
|
4
|
+
from simple_graph_retriever.client import GraphRetrievalClient
|
|
5
|
+
|
|
6
|
+
if __name__ == "__main__":
|
|
7
|
+
print("Starting graph indexing example...")
|
|
8
|
+
client = None
|
|
9
|
+
try:
|
|
10
|
+
# Initialize the client (settings loaded from .env or defaults)
|
|
11
|
+
client = GraphRetrievalClient()
|
|
12
|
+
|
|
13
|
+
# Run the full indexing pipeline
|
|
14
|
+
print(
|
|
15
|
+
"Running client.index()... (This will detect communities, create chunks, and index them)"
|
|
16
|
+
)
|
|
17
|
+
client.index()
|
|
18
|
+
|
|
19
|
+
print("✅ Graph indexing complete.")
|
|
20
|
+
|
|
21
|
+
except Exception as e:
|
|
22
|
+
print(f"❌ An error occurred during indexing: {e}")
|
|
23
|
+
finally:
|
|
24
|
+
if client:
|
|
25
|
+
client.close()
|
|
26
|
+
print("Client closed.")
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
from simple_graph_retriever.client import GraphRetrievalClient
|
|
2
|
+
from simple_graph_retriever.models import RetrievalConfig
|
|
3
|
+
import json
|
|
4
|
+
|
|
5
|
+
if __name__ == "__main__":
|
|
6
|
+
client = GraphRetrievalClient()
|
|
7
|
+
try:
|
|
8
|
+
query = "Pizza"
|
|
9
|
+
|
|
10
|
+
print(f"Retrieving subgraph for query: '{query}'")
|
|
11
|
+
|
|
12
|
+
print("\n--- Full GraphRAG mode ---")
|
|
13
|
+
full_subgraph = client.retriever.retrieve_graph(
|
|
14
|
+
query,
|
|
15
|
+
RetrievalConfig(
|
|
16
|
+
use_communities=True,
|
|
17
|
+
use_chunks=True,
|
|
18
|
+
),
|
|
19
|
+
)
|
|
20
|
+
if full_subgraph:
|
|
21
|
+
print(
|
|
22
|
+
f"Retrieved {len(full_subgraph[0]['nodes'])} nodes and {len(full_subgraph[0]['relationships'])} relationships."
|
|
23
|
+
)
|
|
24
|
+
with open("retrieved_subgraph_full.json", "w") as f:
|
|
25
|
+
json.dump(full_subgraph, f, indent=2)
|
|
26
|
+
else:
|
|
27
|
+
print("No results found.")
|
|
28
|
+
|
|
29
|
+
print("\n--- Chunk-only mode ---")
|
|
30
|
+
chunk_only_subgraph = client.retriever.retrieve_graph(
|
|
31
|
+
query, RetrievalConfig(use_communities=False, use_chunks=True)
|
|
32
|
+
)
|
|
33
|
+
if chunk_only_subgraph:
|
|
34
|
+
print(
|
|
35
|
+
f"Retrieved {len(chunk_only_subgraph[0]['nodes'])} nodes and {len(chunk_only_subgraph[0]['relationships'])} relationships."
|
|
36
|
+
)
|
|
37
|
+
with open("retrieved_subgraph_chunk_only.json", "w") as f:
|
|
38
|
+
json.dump(chunk_only_subgraph, f, indent=2)
|
|
39
|
+
else:
|
|
40
|
+
print("No results found.")
|
|
41
|
+
|
|
42
|
+
print("\n--- Community-only mode ---")
|
|
43
|
+
community_only_subgraph = client.retriever.retrieve_graph(
|
|
44
|
+
query, RetrievalConfig(use_communities=True, use_chunks=False)
|
|
45
|
+
)
|
|
46
|
+
if community_only_subgraph:
|
|
47
|
+
print(
|
|
48
|
+
f"Retrieved {len(community_only_subgraph[0]['nodes'])} nodes and {len(community_only_subgraph[0]['relationships'])} relationships."
|
|
49
|
+
)
|
|
50
|
+
with open("retrieved_subgraph_community_only.json", "w") as f:
|
|
51
|
+
json.dump(community_only_subgraph, f, indent=2)
|
|
52
|
+
else:
|
|
53
|
+
print("No results found.")
|
|
54
|
+
finally:
|
|
55
|
+
client.close()
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "simple-graph-retriever"
|
|
7
|
+
version = "1.0.0rc1"
|
|
8
|
+
authors = [
|
|
9
|
+
{ name="Victor de La Salmonière", email="victor@lettria.com" },
|
|
10
|
+
]
|
|
11
|
+
description = "A simple graph-based retriever using Neo4j and Qdrant"
|
|
12
|
+
readme = "README.md"
|
|
13
|
+
requires-python = ">=3.8"
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"License :: OSI Approved :: MIT License",
|
|
17
|
+
"Operating System :: OS Independent",
|
|
18
|
+
]
|
|
19
|
+
dependencies = [
|
|
20
|
+
"neo4j",
|
|
21
|
+
"qdrant-client",
|
|
22
|
+
"requests",
|
|
23
|
+
"tqdm",
|
|
24
|
+
"igraph",
|
|
25
|
+
"leidenalg",
|
|
26
|
+
"pydantic",
|
|
27
|
+
"pydantic-settings",
|
|
28
|
+
]
|
|
29
|
+
|
|
30
|
+
[tool.hatch.build.targets.wheel]
|
|
31
|
+
packages = ["simple_graph_retriever"]
|
|
File without changes
|
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
from typing import Optional, Tuple
|
|
2
|
+
from .indexer import GraphIndexer
|
|
3
|
+
from .retriever import GraphRetriever
|
|
4
|
+
from .config import settings
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class GraphRetrievalClient:
|
|
8
|
+
"""
|
|
9
|
+
A client for interacting with the graph retrieval API.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
def __init__(self,
|
|
13
|
+
neo4j_uri: Optional[str] = None,
|
|
14
|
+
neo4j_auth: Optional[Tuple[str, str]] = None,
|
|
15
|
+
qdrant_url: Optional[str] = None,
|
|
16
|
+
qdrant_api_key: Optional[str] = None,
|
|
17
|
+
qdrant_chunks_collection: Optional[str] = None,
|
|
18
|
+
qdrant_communities_collection: Optional[str] = None,
|
|
19
|
+
embedder_url: Optional[str] = None,
|
|
20
|
+
vector_size: Optional[int] = None):
|
|
21
|
+
"""
|
|
22
|
+
Initializes the GraphRetrievalClient.
|
|
23
|
+
"""
|
|
24
|
+
neo4j_uri = neo4j_uri or settings.neo4j_uri
|
|
25
|
+
neo4j_auth = neo4j_auth or (settings.neo4j_user, settings.neo4j_password)
|
|
26
|
+
qdrant_url = qdrant_url or settings.qdrant_url
|
|
27
|
+
qdrant_api_key = qdrant_api_key or settings.qdrant_api_key
|
|
28
|
+
qdrant_chunks_collection = qdrant_chunks_collection or settings.qdrant_chunks_collection
|
|
29
|
+
qdrant_communities_collection = qdrant_communities_collection or settings.qdrant_communities_collection
|
|
30
|
+
embedder_url = embedder_url or settings.embedder_url
|
|
31
|
+
vector_size = vector_size or settings.vector_size
|
|
32
|
+
|
|
33
|
+
self._indexer = GraphIndexer(
|
|
34
|
+
neo4j_uri=neo4j_uri,
|
|
35
|
+
neo4j_auth=neo4j_auth,
|
|
36
|
+
qdrant_url=qdrant_url,
|
|
37
|
+
embedder_url=embedder_url,
|
|
38
|
+
qdrant_api_key=qdrant_api_key,
|
|
39
|
+
vector_size=vector_size,
|
|
40
|
+
chunks_collection=qdrant_chunks_collection,
|
|
41
|
+
communities_collection=qdrant_communities_collection
|
|
42
|
+
)
|
|
43
|
+
self._retriever = GraphRetriever(
|
|
44
|
+
qdrant_url=qdrant_url,
|
|
45
|
+
neo4j_uri=neo4j_uri,
|
|
46
|
+
neo4j_auth=neo4j_auth,
|
|
47
|
+
embedder_url=embedder_url,
|
|
48
|
+
qdrant_api_key=qdrant_api_key,
|
|
49
|
+
chunks_collection=qdrant_chunks_collection,
|
|
50
|
+
communities_collection=qdrant_communities_collection
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
@property
|
|
54
|
+
def indexer(self) -> GraphIndexer:
|
|
55
|
+
return self._indexer
|
|
56
|
+
|
|
57
|
+
@property
|
|
58
|
+
def retriever(self) -> GraphRetriever:
|
|
59
|
+
return self._retriever
|
|
60
|
+
|
|
61
|
+
def close(self):
|
|
62
|
+
self._indexer.close()
|
|
63
|
+
self._retriever.close()
|
|
64
|
+
|
|
65
|
+
def index(self):
|
|
66
|
+
"""
|
|
67
|
+
Runs the complete graph indexing pipeline.
|
|
68
|
+
"""
|
|
69
|
+
self.indexer.run_community_detection()
|
|
70
|
+
self.indexer.create_chunks()
|
|
71
|
+
self.indexer.index_chunks()
|
|
72
|
+
self.indexer.index_communities()
|
|
73
|
+
|
|
74
|
+
def clear_index(self):
|
|
75
|
+
"""
|
|
76
|
+
Clears the Qdrant collections.
|
|
77
|
+
"""
|
|
78
|
+
self.indexer.clear_qdrant_collections()
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import os
|
|
3
|
+
import sys
|
|
4
|
+
from typing import Optional
|
|
5
|
+
from pydantic import ValidationError
|
|
6
|
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
7
|
+
|
|
8
|
+
logging.basicConfig(level=os.getenv("LOGLEVEL", "WARNING"))
|
|
9
|
+
logger = logging.getLogger(__name__)
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class Settings(BaseSettings):
|
|
13
|
+
neo4j_uri: str = "bolt://localhost:7687"
|
|
14
|
+
neo4j_user: str = "neo4j"
|
|
15
|
+
neo4j_password: str = "password"
|
|
16
|
+
qdrant_url: str = "http://localhost:6333"
|
|
17
|
+
qdrant_api_key: Optional[str] = None
|
|
18
|
+
qdrant_chunks_collection: str = "chunks"
|
|
19
|
+
qdrant_communities_collection: str = "communities"
|
|
20
|
+
embedder_url: str = "http://localhost:8080"
|
|
21
|
+
vector_size: int = 384
|
|
22
|
+
loglevel: str = "WARNING"
|
|
23
|
+
|
|
24
|
+
model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
try:
|
|
28
|
+
settings = Settings()
|
|
29
|
+
except ValidationError as e:
|
|
30
|
+
missing_fields = [
|
|
31
|
+
error["loc"][0] for error in e.errors() if error["type"] == "missing"
|
|
32
|
+
]
|
|
33
|
+
logger.error(
|
|
34
|
+
f"Missing configuration fields: {', '.join(str(field) for field in missing_fields)}"
|
|
35
|
+
)
|
|
36
|
+
sys.exit(1)
|
|
@@ -0,0 +1,367 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import time
|
|
3
|
+
from typing import List, Dict, Any, Optional
|
|
4
|
+
from neo4j import GraphDatabase
|
|
5
|
+
from qdrant_client import QdrantClient
|
|
6
|
+
from qdrant_client.http.models import PointStruct, Distance, VectorParams
|
|
7
|
+
import requests
|
|
8
|
+
from tqdm import tqdm
|
|
9
|
+
import igraph as ig
|
|
10
|
+
import leidenalg as la
|
|
11
|
+
from .config import logger
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class GraphIndexer:
|
|
15
|
+
def __init__(
|
|
16
|
+
self,
|
|
17
|
+
neo4j_uri,
|
|
18
|
+
neo4j_auth,
|
|
19
|
+
qdrant_url,
|
|
20
|
+
embedder_url,
|
|
21
|
+
qdrant_api_key: Optional[str] = None,
|
|
22
|
+
vector_size=384,
|
|
23
|
+
chunks_collection: str = "chunks",
|
|
24
|
+
communities_collection: str = "communities",
|
|
25
|
+
):
|
|
26
|
+
self.driver = GraphDatabase.driver(neo4j_uri, auth=neo4j_auth)
|
|
27
|
+
self.qdrant = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
|
|
28
|
+
self.embedder_url = embedder_url
|
|
29
|
+
self.vector_batch_size = 100
|
|
30
|
+
self.vector_size = vector_size
|
|
31
|
+
self.chunks_collection = chunks_collection
|
|
32
|
+
self.communities_collection = communities_collection
|
|
33
|
+
self._ensure_qdrant_collections_exist()
|
|
34
|
+
|
|
35
|
+
def _ensure_qdrant_collections_exist(self):
|
|
36
|
+
logger.info("Ensuring Qdrant collections exist...")
|
|
37
|
+
existing_collections = [
|
|
38
|
+
c.name for c in self.qdrant.get_collections().collections
|
|
39
|
+
]
|
|
40
|
+
|
|
41
|
+
if self.chunks_collection not in existing_collections:
|
|
42
|
+
self.qdrant.create_collection(
|
|
43
|
+
collection_name=self.chunks_collection,
|
|
44
|
+
vectors_config=VectorParams(
|
|
45
|
+
size=self.vector_size, distance=Distance.COSINE
|
|
46
|
+
),
|
|
47
|
+
)
|
|
48
|
+
logger.info(f"Created '{self.chunks_collection}' collection.")
|
|
49
|
+
else:
|
|
50
|
+
logger.info(f"'{self.chunks_collection}' collection already exists.")
|
|
51
|
+
|
|
52
|
+
if self.communities_collection not in existing_collections:
|
|
53
|
+
self.qdrant.create_collection(
|
|
54
|
+
collection_name=self.communities_collection,
|
|
55
|
+
vectors_config=VectorParams(
|
|
56
|
+
size=self.vector_size, distance=Distance.COSINE
|
|
57
|
+
),
|
|
58
|
+
)
|
|
59
|
+
logger.info(f"Created '{self.communities_collection}' collection.")
|
|
60
|
+
else:
|
|
61
|
+
logger.info(f"'{self.communities_collection}' collection already exists.")
|
|
62
|
+
logger.info("✅ Qdrant collections checked/created.")
|
|
63
|
+
|
|
64
|
+
def close(self):
|
|
65
|
+
self.driver.close()
|
|
66
|
+
|
|
67
|
+
def clear_qdrant_collections(self):
|
|
68
|
+
logger.info("🗑️ Clearing Qdrant collections...")
|
|
69
|
+
self.qdrant.recreate_collection(
|
|
70
|
+
collection_name=self.chunks_collection,
|
|
71
|
+
vectors_config=VectorParams(
|
|
72
|
+
size=self.vector_size, distance=Distance.COSINE
|
|
73
|
+
),
|
|
74
|
+
)
|
|
75
|
+
self.qdrant.recreate_collection(
|
|
76
|
+
collection_name=self.communities_collection,
|
|
77
|
+
vectors_config=VectorParams(
|
|
78
|
+
size=self.vector_size, distance=Distance.COSINE
|
|
79
|
+
),
|
|
80
|
+
)
|
|
81
|
+
logger.info("✅ Qdrant collections cleared and recreated.")
|
|
82
|
+
|
|
83
|
+
def embed_batch(self, texts: List[str]) -> List[List[float]]:
|
|
84
|
+
"""
|
|
85
|
+
Batch embedding wrapper.
|
|
86
|
+
Adjust payload/response handling based on your specific API.
|
|
87
|
+
"""
|
|
88
|
+
# Example assumes an OpenAI-compatible interface or similar list-in/list-out
|
|
89
|
+
embeddings = []
|
|
90
|
+
for text in texts:
|
|
91
|
+
try:
|
|
92
|
+
resp = requests.post(
|
|
93
|
+
self.embedder_url,
|
|
94
|
+
json={"inputs": [text]}, # Sending as a list of one string
|
|
95
|
+
timeout=30,
|
|
96
|
+
)
|
|
97
|
+
resp.raise_for_status()
|
|
98
|
+
# Assuming response is a list of embeddings, even for single input
|
|
99
|
+
# e.g., [[0.1, 0.2, ...]]
|
|
100
|
+
embeddings.extend(
|
|
101
|
+
resp.json()
|
|
102
|
+
) # Use extend because response is list of lists
|
|
103
|
+
except Exception as e:
|
|
104
|
+
logger.error(f"Embedding failed for text '{text[:50]}...': {e}")
|
|
105
|
+
embeddings.append([]) # Append empty list for failed embeddings
|
|
106
|
+
return embeddings
|
|
107
|
+
|
|
108
|
+
def run_community_detection(self):
|
|
109
|
+
logger.info("1️⃣ Refreshing Community Structure...")
|
|
110
|
+
|
|
111
|
+
# Wipe old community data from Qdrant
|
|
112
|
+
logger.info(f" Wiping Qdrant collection: '{self.communities_collection}'...")
|
|
113
|
+
self.qdrant.recreate_collection(
|
|
114
|
+
collection_name=self.communities_collection,
|
|
115
|
+
vectors_config=VectorParams(
|
|
116
|
+
size=self.vector_size, distance=Distance.COSINE
|
|
117
|
+
),
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
with self.driver.session() as session:
|
|
121
|
+
# Wipe old community data from Neo4j
|
|
122
|
+
logger.info(" Wiping community data from Neo4j...")
|
|
123
|
+
session.run("MATCH (c:Community) DETACH DELETE c")
|
|
124
|
+
session.run("MATCH (n) WHERE n.community_id IS NOT NULL REMOVE n.community_id")
|
|
125
|
+
logger.info(" Community data wiped.")
|
|
126
|
+
|
|
127
|
+
# Fetch all nodes and relationships from Neo4j (excluding chunks and communities)
|
|
128
|
+
logger.info(" Fetching graph data from Neo4j (excluding chunks and communities)...")
|
|
129
|
+
nodes_data = session.run("MATCH (n) WHERE NOT n:GraphChunk AND NOT n:Community RETURN elementId(n) as id").data()
|
|
130
|
+
rels_data = session.run(
|
|
131
|
+
"MATCH (a)-[r]->(b) WHERE NOT a:GraphChunk AND NOT a:Community AND NOT b:GraphChunk AND NOT b:Community RETURN elementId(r) as id, elementId(a) as source, elementId(b) as target"
|
|
132
|
+
).data()
|
|
133
|
+
|
|
134
|
+
node_id_to_idx = {node["id"]: i for i, node in enumerate(nodes_data)}
|
|
135
|
+
|
|
136
|
+
# Create an igraph graph
|
|
137
|
+
g = ig.Graph(directed=True)
|
|
138
|
+
g.add_vertices(len(nodes_data))
|
|
139
|
+
g.vs["neo4j_id"] = [node["id"] for node in nodes_data]
|
|
140
|
+
|
|
141
|
+
edges = []
|
|
142
|
+
for rel in rels_data:
|
|
143
|
+
source_idx = node_id_to_idx.get(rel["source"])
|
|
144
|
+
target_idx = node_id_to_idx.get(rel["target"])
|
|
145
|
+
if source_idx is not None and target_idx is not None:
|
|
146
|
+
edges.append((source_idx, target_idx))
|
|
147
|
+
g.add_edges(edges)
|
|
148
|
+
|
|
149
|
+
logger.info(
|
|
150
|
+
f" Created igraph with {g.vcount()} vertices and {g.ecount()} edges."
|
|
151
|
+
)
|
|
152
|
+
|
|
153
|
+
# Run Leiden algorithm
|
|
154
|
+
logger.info(" Running Leiden algorithm...")
|
|
155
|
+
partition = la.find_partition(g, la.ModularityVertexPartition)
|
|
156
|
+
logger.info(f"✅ Detected {len(partition)} communities.")
|
|
157
|
+
|
|
158
|
+
# Write community IDs back to Neo4j
|
|
159
|
+
logger.info(" Writing community IDs to Neo4j...")
|
|
160
|
+
tx = session.begin_transaction()
|
|
161
|
+
batch_size = 1000
|
|
162
|
+
count = 0
|
|
163
|
+
for i, community_id in enumerate(partition.membership):
|
|
164
|
+
neo4j_node_id = g.vs[i]["neo4j_id"]
|
|
165
|
+
tx.run(
|
|
166
|
+
"MATCH (n) WHERE elementId(n) = $neo4j_node_id SET n.community_id = $community_id",
|
|
167
|
+
neo4j_node_id=neo4j_node_id,
|
|
168
|
+
community_id=community_id,
|
|
169
|
+
)
|
|
170
|
+
count += 1
|
|
171
|
+
if count % batch_size == 0:
|
|
172
|
+
tx.commit()
|
|
173
|
+
tx = session.begin_transaction()
|
|
174
|
+
tx.commit()
|
|
175
|
+
logger.info(f"✅ Wrote community IDs for {count} nodes.")
|
|
176
|
+
|
|
177
|
+
# Materialize Community Nodes (Optional but good for linkage)
|
|
178
|
+
session.run(
|
|
179
|
+
"""
|
|
180
|
+
MATCH (n) WHERE n.community_id IS NOT NULL
|
|
181
|
+
WITH n.community_id AS cid, count(n) as size
|
|
182
|
+
MERGE (c:Community {id: cid})
|
|
183
|
+
SET c.size = size
|
|
184
|
+
"""
|
|
185
|
+
)
|
|
186
|
+
logger.info("✅ Materialized :Community nodes.")
|
|
187
|
+
|
|
188
|
+
def create_chunks(self):
|
|
189
|
+
logger.info("2️⃣ Creating GraphChunks (Node + Context).")
|
|
190
|
+
query = """
|
|
191
|
+
MATCH (n)
|
|
192
|
+
WHERE NOT n:GraphChunk AND NOT n:Community
|
|
193
|
+
// Ensure we only create a chunk if one doesn't already exist for the node
|
|
194
|
+
AND NOT (n)<-[:CENTERED_ON]-(:GraphChunk)
|
|
195
|
+
|
|
196
|
+
// This query will now run once and create all necessary chunks
|
|
197
|
+
WITH n
|
|
198
|
+
|
|
199
|
+
// 1. Gather properties
|
|
200
|
+
WITH n,
|
|
201
|
+
apoc.text.join(labels(n), ", ") as lbls,
|
|
202
|
+
apoc.convert.toJson(properties(n)) as props
|
|
203
|
+
|
|
204
|
+
// 2. Gather 1-hop context (Relationships + Neighbor Labels)
|
|
205
|
+
CALL {
|
|
206
|
+
WITH n
|
|
207
|
+
MATCH (n)-[r]-(m)
|
|
208
|
+
WITH type(r) as rel_type, labels(m) as n_labels, m
|
|
209
|
+
RETURN collect(
|
|
210
|
+
rel_type + " -> " + coalesce(m.name, m.title, head(labels(m)), "Node")
|
|
211
|
+
)[0..10] as context_list // Capping context to 10 neighbors
|
|
212
|
+
}
|
|
213
|
+
|
|
214
|
+
// 3. Format Text Blob
|
|
215
|
+
WITH n, lbls, props,
|
|
216
|
+
"Node: " + lbls + "\nProps: " + props + "\nContext:\n" +
|
|
217
|
+
apoc.text.join(context_list, "\n") as chunk_text
|
|
218
|
+
|
|
219
|
+
// 4. Create Chunk Node
|
|
220
|
+
CREATE (c:GraphChunk {
|
|
221
|
+
id: randomUUID(),
|
|
222
|
+
community_id: n.community_id,
|
|
223
|
+
text: chunk_text
|
|
224
|
+
})
|
|
225
|
+
CREATE (c)-[:CENTERED_ON]->(n)
|
|
226
|
+
RETURN count(c) as created_count
|
|
227
|
+
"""
|
|
228
|
+
|
|
229
|
+
with self.driver.session() as session:
|
|
230
|
+
result = session.run(query).single()
|
|
231
|
+
count = result["created_count"] if result else 0
|
|
232
|
+
logger.info(f"✅ Total chunks created: {count}")
|
|
233
|
+
|
|
234
|
+
def index_chunks(self):
|
|
235
|
+
logger.info("3️⃣ Indexing Chunks to Qdrant...")
|
|
236
|
+
|
|
237
|
+
# Read chunks that are not yet indexed
|
|
238
|
+
fetch_query = """
|
|
239
|
+
MATCH (c:GraphChunk)-[:CENTERED_ON]->(n)
|
|
240
|
+
WHERE c.indexed IS NULL
|
|
241
|
+
RETURN c.id as id, c.text as text, c.community_id as comm_id,
|
|
242
|
+
elementId(n) as center_node_id
|
|
243
|
+
LIMIT $batch_size
|
|
244
|
+
"""
|
|
245
|
+
|
|
246
|
+
mark_done_query = """
|
|
247
|
+
MATCH (c:GraphChunk) WHERE c.id IN $ids
|
|
248
|
+
SET c.indexed = true
|
|
249
|
+
"""
|
|
250
|
+
|
|
251
|
+
with self.driver.session() as session:
|
|
252
|
+
while True:
|
|
253
|
+
# 1. Fetch Batch from Neo4j
|
|
254
|
+
records = session.run(
|
|
255
|
+
fetch_query, batch_size=self.vector_batch_size
|
|
256
|
+
).data()
|
|
257
|
+
if not records:
|
|
258
|
+
break
|
|
259
|
+
|
|
260
|
+
texts = [r["text"] for r in records]
|
|
261
|
+
ids = [r["id"] for r in records]
|
|
262
|
+
|
|
263
|
+
# 2. Embed
|
|
264
|
+
vectors = self.embed_batch(texts)
|
|
265
|
+
|
|
266
|
+
# 3. Prepare Points
|
|
267
|
+
points = []
|
|
268
|
+
for i, rec in enumerate(records):
|
|
269
|
+
if not vectors[i]:
|
|
270
|
+
logger.warning(
|
|
271
|
+
f"Skipping chunk {rec['id']} due to failed embedding."
|
|
272
|
+
)
|
|
273
|
+
continue # skip failed embeddings
|
|
274
|
+
|
|
275
|
+
points.append(
|
|
276
|
+
PointStruct(
|
|
277
|
+
id=rec["id"], # Using UUID from Neo4j
|
|
278
|
+
vector=vectors[i],
|
|
279
|
+
payload={
|
|
280
|
+
"text": rec["text"],
|
|
281
|
+
"community_id": rec["comm_id"],
|
|
282
|
+
"center_node_id": rec[
|
|
283
|
+
"center_node_id"
|
|
284
|
+
], # Store internal ID for retrieval
|
|
285
|
+
"type": "chunk",
|
|
286
|
+
},
|
|
287
|
+
)
|
|
288
|
+
)
|
|
289
|
+
|
|
290
|
+
# 4. Upsert to Qdrant
|
|
291
|
+
if points:
|
|
292
|
+
self.qdrant.upsert(
|
|
293
|
+
collection_name=self.chunks_collection, points=points
|
|
294
|
+
)
|
|
295
|
+
|
|
296
|
+
# 5. Mark as done in Neo4j
|
|
297
|
+
session.run(mark_done_query, ids=ids)
|
|
298
|
+
logger.info(f" Indexed {len(points)} chunks.")
|
|
299
|
+
|
|
300
|
+
logger.info("✅ Chunk indexing complete.")
|
|
301
|
+
|
|
302
|
+
def index_communities(self):
|
|
303
|
+
logger.info("4️⃣ Indexing Communities...")
|
|
304
|
+
|
|
305
|
+
# Strategy: Aggregate the text of the 5 most connected nodes in the community
|
|
306
|
+
query = """
|
|
307
|
+
MATCH (c:Community)
|
|
308
|
+
WHERE c.indexed IS NULL
|
|
309
|
+
|
|
310
|
+
CALL {
|
|
311
|
+
WITH c
|
|
312
|
+
MATCH (n) WHERE n.community_id = c.id
|
|
313
|
+
// Heuristic: Pick 'important' nodes by degree
|
|
314
|
+
WITH n ORDER BY COUNT { (n)--() } DESC LIMIT 5
|
|
315
|
+
|
|
316
|
+
// Re-generate basic text for these nodes
|
|
317
|
+
WITH n, labels(n)[0] + ": " + coalesce(n.name, n.title, "Item") as summary
|
|
318
|
+
RETURN collect(summary) as summaries
|
|
319
|
+
}
|
|
320
|
+
|
|
321
|
+
WITH c, "Community ID: " + toString(c.id) + "\nKey Elements:\n" + apoc.text.join(summaries, "\n") as comm_text
|
|
322
|
+
RETURN c.id as id, comm_text as text
|
|
323
|
+
LIMIT $batch_size
|
|
324
|
+
"""
|
|
325
|
+
|
|
326
|
+
mark_done = "MATCH (c:Community {id: $id}) SET c.indexed = true"
|
|
327
|
+
|
|
328
|
+
with self.driver.session() as session:
|
|
329
|
+
while True:
|
|
330
|
+
records = session.run(query, batch_size=self.vector_batch_size).data()
|
|
331
|
+
if not records:
|
|
332
|
+
break
|
|
333
|
+
|
|
334
|
+
texts = [r["text"] for r in records]
|
|
335
|
+
vectors = self.embed_batch(texts)
|
|
336
|
+
|
|
337
|
+
points = []
|
|
338
|
+
for i, rec in enumerate(records):
|
|
339
|
+
if not vectors[i]:
|
|
340
|
+
continue
|
|
341
|
+
|
|
342
|
+
# Qdrant requires integer or UUID for ID.
|
|
343
|
+
# If community_id is integer, use it directly.
|
|
344
|
+
# If it's a string/uuid, you might need to hash it to UUID.
|
|
345
|
+
points.append(
|
|
346
|
+
PointStruct(
|
|
347
|
+
id=rec["id"],
|
|
348
|
+
vector=vectors[i],
|
|
349
|
+
payload={
|
|
350
|
+
"text": rec["text"],
|
|
351
|
+
"community_id": rec["id"],
|
|
352
|
+
"type": "community",
|
|
353
|
+
},
|
|
354
|
+
)
|
|
355
|
+
)
|
|
356
|
+
|
|
357
|
+
if points:
|
|
358
|
+
self.qdrant.upsert(
|
|
359
|
+
collection_name=self.communities_collection, points=points
|
|
360
|
+
)
|
|
361
|
+
|
|
362
|
+
for r in records:
|
|
363
|
+
session.run(mark_done, id=r["id"])
|
|
364
|
+
|
|
365
|
+
logger.info(f" Indexed {len(points)} communities.")
|
|
366
|
+
|
|
367
|
+
logger.info("✅ Community indexing complete.")
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
from typing import List, Optional
|
|
2
|
+
from pydantic import BaseModel
|
|
3
|
+
|
|
4
|
+
class RetrievalConfig(BaseModel):
|
|
5
|
+
use_communities: bool = True
|
|
6
|
+
use_chunks: bool = True
|
|
7
|
+
max_communities: int = 3
|
|
8
|
+
max_chunks: int = 10
|
|
9
|
+
max_hops: int = 1
|
|
10
|
+
community_score_threshold: float = 0.4
|
|
11
|
+
chunk_score_threshold: float = 0.4
|
|
12
|
+
community_score_drop_off_pct: float = 0.07
|
|
13
|
+
chunk_score_drop_off_pct: float = 0.2
|
|
14
|
+
community_expansion_limit: int = 10
|
|
15
|
+
allowed_rel_types: Optional[List[str]] = None
|
|
16
|
+
denied_rel_types: Optional[List[str]] = None
|
|
@@ -0,0 +1,244 @@
|
|
|
1
|
+
from typing import List, Optional
|
|
2
|
+
import requests
|
|
3
|
+
from qdrant_client import QdrantClient
|
|
4
|
+
from qdrant_client.http.models import Filter, FieldCondition, MatchAny
|
|
5
|
+
from neo4j import GraphDatabase
|
|
6
|
+
|
|
7
|
+
from .models import RetrievalConfig
|
|
8
|
+
from .config import logger
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class GraphRetriever:
|
|
12
|
+
def __init__(
|
|
13
|
+
self,
|
|
14
|
+
qdrant_url,
|
|
15
|
+
neo4j_uri,
|
|
16
|
+
neo4j_auth,
|
|
17
|
+
embedder_url,
|
|
18
|
+
qdrant_api_key: Optional[str] = None,
|
|
19
|
+
chunks_collection: str = "chunks",
|
|
20
|
+
communities_collection: str = "communities",
|
|
21
|
+
):
|
|
22
|
+
self.qdrant_client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
|
|
23
|
+
self.neo4j_driver = GraphDatabase.driver(neo4j_uri, auth=neo4j_auth)
|
|
24
|
+
self.embedder_url = embedder_url
|
|
25
|
+
self.chunks_collection = chunks_collection
|
|
26
|
+
self.communities_collection = communities_collection
|
|
27
|
+
|
|
28
|
+
def close(self):
|
|
29
|
+
self.neo4j_driver.close()
|
|
30
|
+
|
|
31
|
+
def embed(self, text: str) -> list[float]:
|
|
32
|
+
"""
|
|
33
|
+
Embeds a single string using the TEI embedder.
|
|
34
|
+
"""
|
|
35
|
+
try:
|
|
36
|
+
resp = requests.post(self.embedder_url, json={"inputs": text}, timeout=5)
|
|
37
|
+
resp.raise_for_status()
|
|
38
|
+
return resp.json()[0]
|
|
39
|
+
except Exception as e:
|
|
40
|
+
logger.error(f"Embedding failed: {e}")
|
|
41
|
+
return []
|
|
42
|
+
|
|
43
|
+
def search_qdrant(
|
|
44
|
+
self,
|
|
45
|
+
collection: str,
|
|
46
|
+
vector: list[float],
|
|
47
|
+
limit: int,
|
|
48
|
+
filters=None,
|
|
49
|
+
score_threshold=None,
|
|
50
|
+
score_drop_off_pct=None,
|
|
51
|
+
):
|
|
52
|
+
try:
|
|
53
|
+
search_result = self.qdrant_client.search(
|
|
54
|
+
collection_name=collection,
|
|
55
|
+
query_vector=vector,
|
|
56
|
+
limit=limit,
|
|
57
|
+
query_filter=filters,
|
|
58
|
+
score_threshold=score_threshold,
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
# Apply score drop-off filter
|
|
62
|
+
if score_drop_off_pct is not None and len(search_result) > 1:
|
|
63
|
+
filtered_results = [search_result[0]]
|
|
64
|
+
for i in range(1, len(search_result)):
|
|
65
|
+
prev_score = search_result[i - 1].score
|
|
66
|
+
current_score = search_result[i].score
|
|
67
|
+
|
|
68
|
+
if prev_score == 0:
|
|
69
|
+
break
|
|
70
|
+
|
|
71
|
+
drop_off = (prev_score - current_score) / prev_score
|
|
72
|
+
if drop_off > score_drop_off_pct:
|
|
73
|
+
break # Stop including results after a significant drop-off
|
|
74
|
+
filtered_results.append(search_result[i])
|
|
75
|
+
return filtered_results
|
|
76
|
+
|
|
77
|
+
return search_result
|
|
78
|
+
except Exception as e:
|
|
79
|
+
logger.error(f"Qdrant search failed: {e}")
|
|
80
|
+
return []
|
|
81
|
+
|
|
82
|
+
def retrieve_communities(self, query_embedding, config: RetrievalConfig):
|
|
83
|
+
if not config.use_communities:
|
|
84
|
+
return []
|
|
85
|
+
|
|
86
|
+
results = self.search_qdrant(
|
|
87
|
+
collection=self.communities_collection,
|
|
88
|
+
vector=query_embedding,
|
|
89
|
+
limit=config.max_communities,
|
|
90
|
+
score_threshold=config.community_score_threshold,
|
|
91
|
+
score_drop_off_pct=config.community_score_drop_off_pct,
|
|
92
|
+
)
|
|
93
|
+
return [
|
|
94
|
+
{"community_id": r.payload["community_id"], "score": r.score}
|
|
95
|
+
for r in results
|
|
96
|
+
]
|
|
97
|
+
|
|
98
|
+
def retrieve_chunks(
|
|
99
|
+
self, query_embedding, community_results: list, config: RetrievalConfig
|
|
100
|
+
):
|
|
101
|
+
if not config.use_chunks:
|
|
102
|
+
return []
|
|
103
|
+
|
|
104
|
+
filters = None
|
|
105
|
+
if community_results:
|
|
106
|
+
community_ids = [c["community_id"] for c in community_results]
|
|
107
|
+
filters = Filter(
|
|
108
|
+
must=[
|
|
109
|
+
FieldCondition(
|
|
110
|
+
key="community_id", match=MatchAny(any=community_ids)
|
|
111
|
+
)
|
|
112
|
+
]
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
results = self.search_qdrant(
|
|
116
|
+
collection=self.chunks_collection,
|
|
117
|
+
vector=query_embedding,
|
|
118
|
+
limit=config.max_chunks,
|
|
119
|
+
filters=filters,
|
|
120
|
+
score_threshold=config.chunk_score_threshold,
|
|
121
|
+
score_drop_off_pct=config.chunk_score_drop_off_pct,
|
|
122
|
+
)
|
|
123
|
+
|
|
124
|
+
return [
|
|
125
|
+
{"center_node_id": r.payload["center_node_id"], "score": r.score}
|
|
126
|
+
for r in results
|
|
127
|
+
]
|
|
128
|
+
|
|
129
|
+
def retrieve_nodes_from_communities(
|
|
130
|
+
self, community_ids: list, community_expansion_limit: int
|
|
131
|
+
):
|
|
132
|
+
query = """
|
|
133
|
+
MATCH (n)
|
|
134
|
+
WHERE n.community_id IN $cids
|
|
135
|
+
RETURN elementId(n) as id
|
|
136
|
+
LIMIT $limit
|
|
137
|
+
"""
|
|
138
|
+
with self.neo4j_driver.session() as session:
|
|
139
|
+
result = session.run(
|
|
140
|
+
query, cids=community_ids, limit=community_expansion_limit
|
|
141
|
+
)
|
|
142
|
+
return [record["id"] for record in result]
|
|
143
|
+
|
|
144
|
+
def fetch_subgraph(
|
|
145
|
+
self,
|
|
146
|
+
center_node_ids: list,
|
|
147
|
+
max_hops: int,
|
|
148
|
+
allowed_rel_types: Optional[List[str]] = None,
|
|
149
|
+
denied_rel_types: Optional[List[str]] = None,
|
|
150
|
+
):
|
|
151
|
+
rel_type_filter_clauses = []
|
|
152
|
+
if allowed_rel_types:
|
|
153
|
+
rel_type_filter_clauses.append("type(r) IN $allowed_rel_types")
|
|
154
|
+
if denied_rel_types:
|
|
155
|
+
rel_type_filter_clauses.append("NOT type(r) IN $denied_rel_types")
|
|
156
|
+
|
|
157
|
+
rel_filter_string = ""
|
|
158
|
+
if rel_type_filter_clauses:
|
|
159
|
+
rel_filter_string = " WHERE " + " AND ".join(rel_type_filter_clauses)
|
|
160
|
+
|
|
161
|
+
query = f"""
|
|
162
|
+
MATCH (n)
|
|
163
|
+
WHERE elementId(n) IN $ids
|
|
164
|
+
MATCH p=(n)-[r*1..{max_hops}]-(m)
|
|
165
|
+
{rel_filter_string}
|
|
166
|
+
UNWIND nodes(p) as node
|
|
167
|
+
UNWIND relationships(p) as rel
|
|
168
|
+
RETURN collect(DISTINCT node {{.*, element_id: elementId(node)}}) as nodes,
|
|
169
|
+
collect(DISTINCT rel {{.*, element_id: elementId(rel), type: type(rel), start_node_element_id: elementId(startNode(rel)), end_node_element_id: elementId(endNode(rel))}}) as relationships
|
|
170
|
+
"""
|
|
171
|
+
with self.neo4j_driver.session() as session:
|
|
172
|
+
params = {"ids": center_node_ids}
|
|
173
|
+
if allowed_rel_types:
|
|
174
|
+
params["allowed_rel_types"] = allowed_rel_types
|
|
175
|
+
if denied_rel_types:
|
|
176
|
+
params["denied_rel_types"] = denied_rel_types
|
|
177
|
+
|
|
178
|
+
result = session.run(query, params)
|
|
179
|
+
return result.data()
|
|
180
|
+
|
|
181
|
+
def retrieve_graph(
|
|
182
|
+
self, query: str, config: RetrievalConfig, include_chunks: bool = False
|
|
183
|
+
):
|
|
184
|
+
"""
|
|
185
|
+
Retrieves a subgraph from the graph database based on a query string.
|
|
186
|
+
"""
|
|
187
|
+
query_embedding = self.embed(query)
|
|
188
|
+
if not query_embedding:
|
|
189
|
+
return []
|
|
190
|
+
|
|
191
|
+
community_results = self.retrieve_communities(query_embedding, config)
|
|
192
|
+
community_scores = {c["community_id"]: c["score"] for c in community_results}
|
|
193
|
+
|
|
194
|
+
chunk_results = []
|
|
195
|
+
center_node_ids = []
|
|
196
|
+
if config.use_chunks:
|
|
197
|
+
chunk_results = self.retrieve_chunks(
|
|
198
|
+
query_embedding, community_results, config
|
|
199
|
+
)
|
|
200
|
+
center_node_ids += [c["center_node_id"] for c in chunk_results]
|
|
201
|
+
if config.use_communities:
|
|
202
|
+
community_ids = [c["community_id"] for c in community_results]
|
|
203
|
+
center_node_ids += self.retrieve_nodes_from_communities(
|
|
204
|
+
community_ids,
|
|
205
|
+
community_expansion_limit=config.community_expansion_limit,
|
|
206
|
+
)
|
|
207
|
+
|
|
208
|
+
if not center_node_ids:
|
|
209
|
+
return []
|
|
210
|
+
# print(f"Number of center nodes: {len(center_node_ids)}")
|
|
211
|
+
subgraph = self.fetch_subgraph(
|
|
212
|
+
center_node_ids,
|
|
213
|
+
config.max_hops,
|
|
214
|
+
config.allowed_rel_types,
|
|
215
|
+
config.denied_rel_types,
|
|
216
|
+
)
|
|
217
|
+
# sub_graph_ids = [node["element_id"] for node in subgraph[0]["nodes"]]
|
|
218
|
+
# print(f"Subgraph node IDs: {sub_graph_ids}")
|
|
219
|
+
if subgraph and subgraph[0].get("nodes"):
|
|
220
|
+
chunk_scores = {c["center_node_id"]: c["score"] for c in chunk_results}
|
|
221
|
+
|
|
222
|
+
for node in subgraph[0]["nodes"]:
|
|
223
|
+
node_id = node.get("element_id")
|
|
224
|
+
community_id = node.get("community_id")
|
|
225
|
+
if node_id in chunk_scores:
|
|
226
|
+
node["chunk_score"] = chunk_scores[node_id]
|
|
227
|
+
if community_id in community_scores:
|
|
228
|
+
node["community_score"] = community_scores[community_id]
|
|
229
|
+
|
|
230
|
+
if not include_chunks:
|
|
231
|
+
nodes_to_keep = [
|
|
232
|
+
node for node in subgraph[0]["nodes"] if "text" not in node
|
|
233
|
+
]
|
|
234
|
+
node_ids_to_keep = {node["element_id"] for node in nodes_to_keep}
|
|
235
|
+
relationships_to_keep = [
|
|
236
|
+
rel
|
|
237
|
+
for rel in subgraph[0]["relationships"]
|
|
238
|
+
if rel["start_node_element_id"] in node_ids_to_keep
|
|
239
|
+
and rel["end_node_element_id"] in node_ids_to_keep
|
|
240
|
+
]
|
|
241
|
+
subgraph[0]["nodes"] = nodes_to_keep
|
|
242
|
+
subgraph[0]["relationships"] = relationships_to_keep
|
|
243
|
+
|
|
244
|
+
return subgraph
|