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.
@@ -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,6 @@
1
+ .env*
2
+ dist/
3
+ *.cql
4
+ *.ttl
5
+ *.pdf
6
+ examples/save_neo4j
@@ -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"]
@@ -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