Skip to main content

Graph-based memory system using DuckDB

Project description

GraphMemory - GraphRAG Database

GraphMemory

An embedded graph database for RAG and knowledge graph applications, powered by DuckDB. Vector similarity search, full-text search, hybrid search, merge/upsert, graph traversal, and a full GraphRAG retrieval pipeline — all in a single Python package.

Features

  • Vector Search — HNSW-indexed nearest neighbors (L2, cosine, inner product)
  • Full-Text Search — BM25-scored search across node properties
  • Hybrid Search — Combined vector + text with configurable weights
  • GraphRAG — Retrieval pipeline: hybrid search → graph expansion → context assembly → LLM Q&A
  • Merge / Upsert — Deduplicate nodes by property keys and edges by (source, target, relation)
  • Query Builder — Fluent, parameterized API with multi-hop traversal
  • DSPy Extraction — Entity/relationship extraction from text via DSPy (optional)
  • Graph Algorithms — PageRank, centrality, components via NetworkX (optional)
  • Import / Export — JSON, CSV, GraphML
  • Visualizer — Interactive D3.js force-directed graph in the browser
  • Thread-Safe — Connection pooling, transactions, automatic retry with exponential backoff

Installation

pip install graphmemory

# Optional
pip install graphmemory[extraction]   # DSPy extraction
pip install graphmemory[algorithms]   # NetworkX algorithms

Quick Start

from graphmemory import GraphMemory, Node, Edge

graph = GraphMemory(database="graph.db", vector_length=3, distance_metric="cosine")

# Insert nodes
alice = Node(type="Person", properties={"name": "Alice", "role": "engineer"}, vector=[0.1, 0.8, 0.3])
bob = Node(type="Person", properties={"name": "Bob", "role": "manager"}, vector=[0.2, 0.7, 0.4])
graph.insert_node(alice)
graph.insert_node(bob)

# Insert edge
graph.insert_edge(Edge(source_id=alice.id, target_id=bob.id, relation="reports_to", weight=1.0))

# Vector search
nearest = graph.nearest_nodes(vector=[0.1, 0.8, 0.3], limit=5)

# Full-text search
results = graph.search_nodes("engineer", limit=10)

# Hybrid search
results = graph.hybrid_search("engineer", query_vector=[0.1, 0.8, 0.3], text_weight=0.5, vector_weight=0.5)

# Context manager
with GraphMemory(database="graph.db", vector_length=3) as graph:
    graph.insert_node(alice)

Usage

Query Builder

# Filter by type and properties
results = graph.query().match(type="Person").where(role="engineer").execute()

# Multi-hop traversal
results = graph.query().traverse(source_id=alice.id, depth=2).execute()

# Paginate and order
results = graph.query().match(type="Person").order_by("name").limit(10).offset(0).execute()

# Query edges
edges = graph.query().match(type="Person").edges().execute()

Merge / Upsert

Insert-or-update nodes matched by property keys. Edges deduplicate on (source_id, target_id, relation).

from graphmemory import MergeStrategy

# Insert if no match, update if "name" matches an existing Person node
result = graph.merge_node(alice, match_keys=["name"])
print(result.created)  # True = inserted, False = updated

# Bulk merge with strategy
results = graph.bulk_merge_nodes(nodes, match_keys=["name"], strategy=MergeStrategy.UPDATE)

# Edge merge
result = graph.merge_edge(edge)
results = graph.bulk_merge_edges(edges)
Strategy Behavior
UPDATE Shallow merge — existing keys preserved, incoming keys added/overwritten (default)
REPLACE Incoming properties fully replace existing
KEEP Existing properties unchanged; only new nodes inserted

GraphRAG Retrieval

Full pipeline: hybrid search → multi-hop graph expansion → token-aware context assembly → LLM generation.

# Retrieve context
result = graph.retrieve(query="Who leads ML?", query_vector=embedding, max_hops=2, max_tokens=4000)
print(result.context_text)      # Prompt-ready string
print(result.token_estimate)    # Token count estimate

# End-to-end Q&A
answer = graph.ask(query="Who leads ML?", query_vector=embedding, llm_callable=my_llm)
print(answer["answer"])

DSPy Extraction

Requires pip install graphmemory[extraction]. Uses DSPy typed predictors to extract entities and relationships from text.

from graphmemory.extraction import extract_and_store, extract_and_merge
import dspy

dspy.configure(lm=dspy.LM("openai/gpt-4o-mini"))

text = """George Washington was the first President. Thomas Jefferson
served as Secretary of State under Washington."""

# Extract and insert (may create duplicates on repeated calls)
nodes, edges = extract_and_store(graph, text)

# Extract and merge (deduplicates against existing graph)
node_results, edge_results = extract_and_merge(graph, text, match_keys=["name"])
Function Description
extract_nodes(text) Extract entity nodes from text
extract_edges(text, nodes) Extract relationships between known nodes
extract(text) Extract both nodes and edges
extract_and_store(graph, text) Extract and insert into graph
extract_and_merge(graph, text, match_keys) Extract and merge (deduplicated)

Graph Algorithms

Requires pip install graphmemory[algorithms]. Powered by NetworkX.

from graphmemory.algorithms import pagerank, betweenness_centrality, connected_components, to_networkx

scores = pagerank(graph)
centrality = betweenness_centrality(graph)
components = connected_components(graph)
G = to_networkx(graph)  # Export to NetworkX DiGraph
Function Description
pagerank(graph, alpha=0.85) PageRank scores for all nodes
betweenness_centrality(graph) Betweenness centrality scores
degree_distribution(graph) In/out/total degree per node
connected_components(graph) Weakly connected components (largest first)
to_networkx(graph) Export to networkx.DiGraph

Import / Export

# Export
data = graph.export_graph(format="json")       # also: "csv", "graphml", "json_string"

# Import
graph.import_graph(data, format="json")

Visualizer

Interactive D3.js force-directed graph visualization — opens in your browser with zero dependencies.

# Open in browser
graph.visualize()

# Save to file
graph.visualize(output="my_graph.html", open_browser=False)

Features: drag nodes, zoom/pan, hover to highlight connections, click for detail panel, search bar, filter by node type.

Data Models

Model Fields
Node id: UUID, type: str, properties: dict, vector: list[float]
Edge id: UUID, source_id: UUID, target_id: UUID, relation: str, weight: float
NearestNode node: Node, distance: float
SearchResult node: Node, score: float
TraversalResult node: Node, depth: int, path: list[UUID]
MergeResult node: Node, created: bool
EdgeMergeResult edge: Edge, created: bool
RetrievalResult query: str, contexts: list, context_text: str, token_estimate: int

All IDs are auto-generated UUIDs. All models are Pydantic BaseModel instances.

API Reference

Connection

Method Description
GraphMemory(database=None, vector_length=3, distance_metric='l2', hnsw_ef_construction=128, hnsw_ef_search=64, hnsw_m=16, auto_index=True) Initialize. None = in-memory. HNSW index auto-created.
close() Close connection (thread-safe, idempotent).
transaction() Context manager for atomic operations.

Nodes

Method Description
insert_node(node) -> UUID Insert a node.
bulk_insert_nodes(nodes) -> list[Node] Bulk insert.
merge_node(node, match_keys, strategy=UPDATE) -> MergeResult Insert or update by property match.
bulk_merge_nodes(nodes, match_keys, ...) -> list[MergeResult] Bulk merge.
get_node(node_id) -> Node Get by ID.
update_node(node_id, **kwargs) -> bool Update fields.
delete_node(node_id) Delete node and its edges.
bulk_delete_nodes(node_ids) Bulk delete.
nodes_by_attribute(attr, value) -> list[Node] Query by property.

Edges

Method Description
insert_edge(edge) Insert an edge.
bulk_insert_edges(edges) Bulk insert.
merge_edge(edge) -> EdgeMergeResult Insert or update by (source, target, relation).
bulk_merge_edges(edges) -> list[EdgeMergeResult] Bulk merge.
get_edge(edge_id) -> Edge Get by ID.
update_edge(edge_id, **kwargs) -> bool Update fields.
delete_edge(source_id, target_id) Delete by endpoints.
bulk_delete_edges(edge_ids) Bulk delete.

Search

Method Description
nearest_nodes(vector, limit) -> list[NearestNode] Vector similarity search.
search_nodes(query_text, limit=10) -> list[SearchResult] Full-text BM25 search.
hybrid_search(query_text, query_vector, ...) -> list[SearchResult] Combined text + vector search.
create_index(ef_construction=None, ef_search=None, m=None) Create/recreate HNSW index with tunable params. Auto-called on init.
compact_index() Compact HNSW index to reclaim space after deletions.

Retrieval

Method Description
retrieve(query, query_vector, ...) -> RetrievalResult Full GraphRAG retrieval pipeline.
ask(query, query_vector, llm_callable, ...) -> dict Retrieval + LLM generation.

Traversal

Method Description
connected_nodes(node_id) -> list[Node] All nodes connected to a node.
query() -> QueryBuilder Fluent query builder.

Import / Export

Method Description
export_graph(format='json') Export as JSON, CSV, GraphML, or JSON string.
import_graph(data, format='json') Import from any supported format.
visualize(output=None, open_browser=True) -> str Interactive D3.js graph visualization in the browser.

Examples

See examples/ for complete usage:

  • openai_example.py — OpenAI embeddings, similarity search, attribute queries
  • lexical_graph.py — Wikipedia text with SentenceTransformer embeddings
  • dspy_example_typed_pred.py — Knowledge graph extraction with DSPy

Testing

296 tests covering all functionality.

python3 -m pytest tests/tests.py -v

License

MIT License. See LICENSE.

Contributing

Contributions welcome — open an issue or submit a PR.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

graphmemory-1.4.0.tar.gz (4.8 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

graphmemory-1.4.0-py3-none-any.whl (33.3 kB view details)

Uploaded Python 3

File details

Details for the file graphmemory-1.4.0.tar.gz.

File metadata

  • Download URL: graphmemory-1.4.0.tar.gz
  • Upload date:
  • Size: 4.8 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.3

File hashes

Hashes for graphmemory-1.4.0.tar.gz
Algorithm Hash digest
SHA256 30afb79ad0a99b2f8a6da3e2ccfc725e509f4cbd20d17ea5ec8703ee8069d704
MD5 9252ce86dc95364a2b88b85ce785e3b3
BLAKE2b-256 4b311dc6488d88060f8861a51e7ead1a8a8f03d5d4c88b67dfa312833e3f3508

See more details on using hashes here.

File details

Details for the file graphmemory-1.4.0-py3-none-any.whl.

File metadata

  • Download URL: graphmemory-1.4.0-py3-none-any.whl
  • Upload date:
  • Size: 33.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.3

File hashes

Hashes for graphmemory-1.4.0-py3-none-any.whl
Algorithm Hash digest
SHA256 8ad034be76ca760ea7304c4290d015f409bfaee8187445d7c93b2369fd19dacd
MD5 ee3b133234fda0bdec7768b77f58b1fa
BLAKE2b-256 ec0dcb0ce75188cb9d02be0476bec6df6c8a46b95b89633a84724875421ae80e

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page