Skip to main content

Graph-based memory system using DuckDB

Project description

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') Initialize. None = in-memory.
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() Create HNSW index for faster vector search.

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

265 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.1.2.tar.gz (76.2 kB view details)

Uploaded Source

Built Distribution

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

graphmemory-1.1.2-py3-none-any.whl (24.8 kB view details)

Uploaded Python 3

File details

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

File metadata

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

File hashes

Hashes for graphmemory-1.1.2.tar.gz
Algorithm Hash digest
SHA256 813a746a02559c65d0c4bda017716832f9663ed5e552e3c1c82e6feddcedc68e
MD5 f7e868249855d47eebbbcd146514ab33
BLAKE2b-256 ef48a8492945973867105b9f1d0083782d0620314efe2c4dc1873942943ed4c7

See more details on using hashes here.

File details

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

File metadata

  • Download URL: graphmemory-1.1.2-py3-none-any.whl
  • Upload date:
  • Size: 24.8 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.1.2-py3-none-any.whl
Algorithm Hash digest
SHA256 97a4a59dea114e9a87746c37d66d7c48dbca2058c16c4232b431e561368fbc26
MD5 4b58673f841deed3ded1a4d8538efed6
BLAKE2b-256 8c3f01f70a6662795d9f1c45b30922435e012deafe1fdbe0ec0548a99b290037

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