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
  • 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")

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.

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

291 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.3.0.tar.gz (4.7 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.3.0-py3-none-any.whl (29.2 kB view details)

Uploaded Python 3

File details

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

File metadata

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

File hashes

Hashes for graphmemory-1.3.0.tar.gz
Algorithm Hash digest
SHA256 041d079da4530c2f5309f0e1e781d84a69c978b1da13bbe8ab04eab897ac60fc
MD5 829d51c911b0431d1b6b944f00a14eaa
BLAKE2b-256 36bae917ead4fca71c1937fdf6637e6b6ededfdd4a798128c0a1175dc767d945

See more details on using hashes here.

File details

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

File metadata

  • Download URL: graphmemory-1.3.0-py3-none-any.whl
  • Upload date:
  • Size: 29.2 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.3.0-py3-none-any.whl
Algorithm Hash digest
SHA256 f60358e7d147092b01d3537de519e125a38d57fdbe21c8c4ec30804de52aaeb7
MD5 5364d93f83b731ad4bab2ed7885d6adc
BLAKE2b-256 d9196304307d8877dcf7010c6fcd18d0e1ca5d3055ca6845af9c551cf217cda1

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