Skip to content

Commit 89abab9

Browse files
KSemenenkoKSemenenko
committed
Add Apache AGE support with new types and integration methods
1 parent d481406 commit 89abab9

33 files changed

+2903
-60
lines changed

‎AGENTS.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
- Register language models through Microsoft.Extensions.AI keyed services; avoid bespoke `LanguageModelConfig` providers.
99
- Always run `dotnet format GraphRag.slnx` before finishing work.
1010
- Always run `dotnet test GraphRag.slnx` before finishing work, after building.
11+
- Avoid per-client connection locks (e.g., `_connectionLock` in `AgeClient`); rely on the smart connection manager for concurrency.
1112

1213
# Conversations
1314
any resulting updates to agents.md should go under the section "## Rules to follow"

‎Directory.Build.props‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@
2525
<RepositoryUrl>https://github.com/managedcode/graphrag</RepositoryUrl>
2626
<PackageProjectUrl>https://github.com/managedcode/graphrag</PackageProjectUrl>
2727
<Product>Managed Code GraphRag</Product>
28-
<Version>0.0.4</Version>
29-
<PackageVersion>0.0.4</PackageVersion>
28+
<Version>0.0.5</Version>
29+
<PackageVersion>0.0.5</PackageVersion>
3030

3131
</PropertyGroup>
3232
<PropertyGroup Condition="'$(GITHUB_ACTIONS)' == 'true'">

‎README.md‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,49 @@ dotnet test tests/ManagedCode.GraphRag.Tests/ManagedCode.GraphRag.Tests.csproj \
205205

206206
---
207207

208+
## Apache AGE / PostgreSQL Setup
209+
210+
GraphRAG ships with a first-class Apache AGE adapter (`ManagedCode.GraphRag.Postgres`). AGE is enabled on top of PostgreSQL, so you only need a standard Postgres instance with the AGE extension installed.
211+
212+
1. **Run an AGE-enabled Postgres instance.** The integration tests use the official container and you can do the same locally:
213+
```bash
214+
docker run --rm \
215+
-e POSTGRES_USER=postgres \
216+
-e POSTGRES_PASSWORD=postgres \
217+
-e POSTGRES_DB=graphrag \
218+
-p 5432:5432 \
219+
apache/age:latest
220+
```
221+
2. **Provide a connection string.** `AgeConnectionManager` accepts a standard Npgsql-style string (for example `Host=localhost;Port=5432;Username=postgres;Password=postgres;Database=graphrag`). The manager automatically runs `CREATE EXTENSION IF NOT EXISTS age;`, `LOAD 'age';`, and `SET search_path = ag_catalog, "$user", public;` before any query executes.
222+
3. **Configure the store.** Either bind `PostgresGraphStoreOptions` in code or use configuration. The snippet below shows the JSON shape (environment variables can follow the same hierarchy, e.g. `GraphRag__GraphStores__postgres__ConnectionString`):
223+
```json
224+
{
225+
"GraphRag": {
226+
"GraphStores": {
227+
"postgres": {
228+
"ConnectionString": "Host=localhost;Port=5432;Username=postgres;Password=postgres;Database=graphrag",
229+
"GraphName": "graphrag",
230+
"MaxConnections": 40,
231+
"MakeDefault": true
232+
}
233+
}
234+
}
235+
}
236+
```
237+
4. **Register through DI.** `services.AddPostgresGraphStore("postgres", configure: ...)` wires up `IAgeConnectionManager`, `IAgeClientFactory`, and `IGraphStore` automatically. `MaxConnections` caps the number of concurrent AGE sessions (the default is 40 so you stay under the container’s `max_connections`).
238+
239+
> **Tip:** `IGraphStore` now exposes `GetNodesAsync` and `GetRelationshipsAsync` in addition to the targeted APIs (`InitializeAsync`, `Upsert*`, `GetOutgoingRelationshipsAsync`). These use the new AGE-powered enumerations so you can inspect or export the full graph without dropping down to concrete implementations.
240+
241+
> **Pagination:** `GetNodesAsync` and `GetRelationshipsAsync` accept an optional `GraphTraversalOptions` object (`new GraphTraversalOptions { Skip = 100, Take = 50 }`) if you want to page through very large graphs. The defaults stream everything, one record at a time, without materialising the entire graph in memory.
242+
243+
---
244+
245+
## Credits
246+
247+
- **pg-age** ([Allison-E/pg-age](https://github.com/Allison-E/pg-age)) — we vendor this Apache AGE client library (see `src/ManagedCode.GraphRag.Postgres/ApacheAge`) so GraphRAG for .NET can rely on a battle-tested connector. Many thanks to Allison and contributors for making AGE on PostgreSQL accessible.
248+
249+
---
250+
208251
## Additional Documentation & Diagrams
209252

210253
- [`docs/indexing-and-query.md`](docs/indexing-and-query.md) explains how each workflow maps to the GraphRAG research diagrams (default data flow, query orchestrations, prompt tuning strategies) published at [microsoft.github.io/graphrag](https://microsoft.github.io/graphrag/).

‎src/ManagedCode.GraphRag.CosmosDb/CosmosGraphStore.cs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,4 +69,72 @@ async IAsyncEnumerable<GraphRelationship> Fetch([EnumeratorCancellation] Cancell
6969
private sealed record NodeDocument(string Id, string Label, Dictionary<string, object?> Properties);
7070

7171
private sealed record EdgeDocument(string Id, string SourceId, string TargetId, string Type, Dictionary<string, object?> Properties);
72+
73+
public IAsyncEnumerable<GraphNode> GetNodesAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
74+
{
75+
options?.Validate();
76+
return Fetch(options, cancellationToken);
77+
78+
async IAsyncEnumerable<GraphNode> Fetch(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
79+
{
80+
var container = _client.GetContainer(_databaseId, _nodesContainerId);
81+
IQueryable<NodeDocument> queryable = container.GetItemLinqQueryable<NodeDocument>(allowSynchronousQueryExecution: false);
82+
queryable = queryable.OrderBy(node => node.Id);
83+
84+
if (traversalOptions?.Skip is > 0 and var skip)
85+
{
86+
queryable = queryable.Skip(skip);
87+
}
88+
89+
if (traversalOptions?.Take is { } take)
90+
{
91+
queryable = queryable.Take(take);
92+
}
93+
94+
using var iterator = queryable.ToFeedIterator();
95+
96+
while (iterator.HasMoreResults)
97+
{
98+
var response = await iterator.ReadNextAsync(token);
99+
foreach (var node in response)
100+
{
101+
yield return new GraphNode(node.Id, node.Label, node.Properties);
102+
}
103+
}
104+
}
105+
}
106+
107+
public IAsyncEnumerable<GraphRelationship> GetRelationshipsAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
108+
{
109+
options?.Validate();
110+
return FetchEdges(options, cancellationToken);
111+
112+
async IAsyncEnumerable<GraphRelationship> FetchEdges(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
113+
{
114+
var container = _client.GetContainer(_databaseId, _edgesContainerId);
115+
IQueryable<EdgeDocument> queryable = container.GetItemLinqQueryable<EdgeDocument>(allowSynchronousQueryExecution: false);
116+
queryable = queryable.OrderBy(edge => edge.SourceId).ThenBy(edge => edge.TargetId);
117+
118+
if (traversalOptions?.Skip is > 0 and var skip)
119+
{
120+
queryable = queryable.Skip(skip);
121+
}
122+
123+
if (traversalOptions?.Take is { } take)
124+
{
125+
queryable = queryable.Take(take);
126+
}
127+
128+
using var iterator = queryable.ToFeedIterator();
129+
130+
while (iterator.HasMoreResults)
131+
{
132+
var response = await iterator.ReadNextAsync(token);
133+
foreach (var edge in response)
134+
{
135+
yield return new GraphRelationship(edge.SourceId, edge.TargetId, edge.Type, edge.Properties);
136+
}
137+
}
138+
}
139+
}
72140
}

‎src/ManagedCode.GraphRag.Neo4j/Neo4jGraphStore.cs‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,88 @@ async IAsyncEnumerable<GraphRelationship> Fetch([EnumeratorCancellation] Cancell
7777
}
7878
}
7979

80+
public IAsyncEnumerable<GraphRelationship> GetRelationshipsAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
81+
{
82+
options?.Validate();
83+
cancellationToken.ThrowIfCancellationRequested();
84+
return FetchAll(options, cancellationToken);
85+
86+
async IAsyncEnumerable<GraphRelationship> FetchAll(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
87+
{
88+
await using var session = _driver.AsyncSession();
89+
var parameters = new Dictionary<string, object>();
90+
var pagination = BuildNeo4jPaginationClause(traversalOptions, parameters, orderBy: "source.id, target.id, type(rel)");
91+
var cursor = await session.RunAsync(
92+
@"MATCH (source)-[rel]->(target)
93+
RETURN source.id AS SourceId, target.id AS TargetId, type(rel) AS Type, properties(rel) AS Properties" + pagination,
94+
parameters);
95+
96+
while (await cursor.FetchAsync())
97+
{
98+
token.ThrowIfCancellationRequested();
99+
var record = cursor.Current;
100+
var properties = record["Properties"].As<IDictionary<string, object?>>();
101+
yield return new GraphRelationship(
102+
record["SourceId"].As<string>(),
103+
record["TargetId"].As<string>(),
104+
record["Type"].As<string>(),
105+
properties.ToDictionary(kvp => kvp.Key, kvp => kvp.Value));
106+
}
107+
}
108+
}
109+
110+
public IAsyncEnumerable<GraphNode> GetNodesAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
111+
{
112+
options?.Validate();
113+
cancellationToken.ThrowIfCancellationRequested();
114+
return FetchNodes(options, cancellationToken);
115+
116+
async IAsyncEnumerable<GraphNode> FetchNodes(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
117+
{
118+
await using var session = _driver.AsyncSession();
119+
var parameters = new Dictionary<string, object>();
120+
var pagination = BuildNeo4jPaginationClause(traversalOptions, parameters, orderBy: "n.id");
121+
var cursor = await session.RunAsync(
122+
@"MATCH (n)
123+
RETURN head(labels(n)) AS Label, n.id AS Id, properties(n) AS Properties" + pagination,
124+
parameters);
125+
126+
while (await cursor.FetchAsync())
127+
{
128+
token.ThrowIfCancellationRequested();
129+
var record = cursor.Current;
130+
var properties = record["Properties"].As<IDictionary<string, object?>>();
131+
yield return new GraphNode(
132+
record["Id"].As<string>(),
133+
record["Label"].As<string>(),
134+
properties.ToDictionary(kvp => kvp.Key, kvp => kvp.Value));
135+
}
136+
}
137+
}
138+
139+
private static string BuildNeo4jPaginationClause(GraphTraversalOptions? options, IDictionary<string, object> parameters, string orderBy)
140+
{
141+
var clause = $" ORDER BY {orderBy}";
142+
if (options is null)
143+
{
144+
return clause;
145+
}
146+
147+
if (options.Skip is > 0 and var skip)
148+
{
149+
parameters["skip"] = skip;
150+
clause += " SKIP $skip";
151+
}
152+
153+
if (options.Take is { } take)
154+
{
155+
parameters["limit"] = take;
156+
clause += " LIMIT $limit";
157+
}
158+
159+
return clause;
160+
}
161+
80162
private static string EscapeLabel(string value)
81163
{
82164
if (value.Any(ch => !char.IsLetterOrDigit(ch) && ch != '_' && ch != ':'))

‎src/ManagedCode.GraphRag.Postgres/ApacheAge/AGEClientEventId.cs‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
namespace GraphRag.Storage.Postgres.ApacheAge;
2+
3+
internal static class AgeClientEventId
4+
{
5+
#region Connection
6+
public const int CONNECTION_OPENED = 1000;
7+
public const int CONNECTION_CLOSED = 1001;
8+
9+
public const int NULL_CONNECTION_WARNING = 1800;
10+
11+
public const int OPEN_CONNECTION_ERROR = 1900;
12+
public const int CLOSE_CONNECTION_ERROR = 1901;
13+
#endregion
14+
15+
#region Internals
16+
public const int EXTENSION_CREATED = 2101;
17+
public const int EXTENSION_LOADED = 2103;
18+
public const int RETRIEVED_CURRENT_SEARCH_PATH = 2104;
19+
public const int AG_CATALOG_ADDED_TO_SEARCH_PATH = 2105;
20+
21+
public const int EXTENSION_NOT_CREATED_ERROR = 2900;
22+
public const int EXTENSION_NOT_LOADED_ERROR = 2901;
23+
public const int AG_CATALOG_NOT_ADDED_TO_SEARCH_PATH_ERROR = 2902;
24+
#endregion
25+
26+
#region Commands
27+
public const int GRAPH_CREATED = 3001;
28+
public const int GRAPH_DROPPED = 3002;
29+
public const int GRAPH_EXISTS = 3003;
30+
public const int GRAPH_DOES_NOT_EXIST = 3004;
31+
32+
public const int CYPHER_EXECUTED = 3101;
33+
public const int QUERY_EXECUTED = 3102;
34+
35+
public const int GRAPH_NOT_CREATED_ERROR = 3900;
36+
public const int GRAPH_NOT_DROPPED_ERROR = 3901;
37+
public const int CYPHER_EXECUTION_ERROR = 3902;
38+
public const int QUERY_EXECUTION_ERROR = 3903;
39+
#endregion
40+
41+
public const int UNKNOWN_ERROR = 9900;
42+
}

0 commit comments

Comments
 (0)