bradagi/graphmemory/graphmemory

pip install graphmemory
pip install graphmemory[extraction]   # DSPy entity/relation extraction
pip install graphmemory[algorithms]   # NetworkX algorithms

from graphmemory import GraphMemory, Node, Edge, MergeStrategy
# database=None is in-memory; pass a path for persistence.
# vector_length and distance_metric are fixed at init time.
graph = GraphMemory(
    database="graph.db",
    vector_length=1536,              # must match your embedding model
    distance_metric="cosine",        # "l2" | "cosine" | "inner_product"
    hnsw_ef_construction=128,
    hnsw_ef_search=64,
    hnsw_m=16,
    auto_index=True,                 # HNSW auto-built on init
    max_retries=3,                   # transient IO error retry
)

alice = Node(type="Person", properties={"name": "Alice"}, vector=embed("Alice"))
bob = Node(type="Person", properties={"name": "Bob"}, vector=embed("Bob"))
graph.insert_node(alice)
graph.insert_node(bob)
graph.insert_edge(Edge(source_id=alice.id, target_id=bob.id, relation="reports_to"))
# Idempotent re-ingest on a natural key
graph.merge_node(alice, match_keys=["name"])
# Fuzzy merge — tolerates "Alice Smith" vs "alice smith"
graph.merge_node(
    alice,
    match_keys=["name"],
    similarity_threshold=0.9,        # Jaro-Winkler threshold (1.0 = exact)
    vector_threshold=0.2,            # optional cosine distance cap
    match_type=True,                 # also require same `type`
    strategy=MergeStrategy.UPDATE,   # UPDATE | REPLACE | KEEP
)

results = graph.hybrid_search(
    query_text="who leads ML?",
    query_vector=embed("who leads ML?"),
    text_weight=0.5,
    vector_weight=0.5,
    limit=10,
)
for r in results:
    print(r.score, r.node.properties)

# Context-only (own the prompt)
result = graph.retrieve(
    query=q, query_vector=qv,
    max_hops=2, max_tokens=4000, search_limit=10,
)
print(result.context_text, result.token_estimate, result.seed_node_count, result.total_node_count)
# End-to-end — llm_callable signature: (system_prompt, user_prompt) -> str
def my_llm(system, user):
    return openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "system", "content": system}, {"role": "user", "content": user}],
    ).choices[0].message.content
answer = graph.ask(query=q, query_vector=qv, llm_callable=my_llm)
print(answer["answer"])

# Filter by type + property
engineers = graph.query().match(type="Person").where(role="engineer").execute()
# Multi-hop traversal — returns TraversalResult with depth + path
two_hop = graph.query().traverse(source_id=alice.id, depth=2).execute()
# Paginate + order
page = graph.query().match(type="Person").order_by("name").limit(20).offset(40).execute()
# Return edges instead of nodes
edges = graph.query().match(type="Person").edges().execute()

import dspy
from graphmemory.extraction import extract_and_merge, extract_and_merge_parallel
dspy.configure(lm=dspy.LM("openai/gpt-4o-mini"))
# Single pass
node_results, edge_results = extract_and_merge(
    graph, text, match_keys=["name"], similarity_threshold=0.88,
)
# Parallel across chunks — two phases: nodes first (all chunks), then edges
# with the full node context. Saturates your RPM.
node_results, edge_results = extract_and_merge_parallel(
    graph,
    chunks=paragraph_chunks,
    match_keys=["name"],
    similarity_threshold=0.88,
    max_workers=8,                   # match your provider's RPM headroom
    on_progress=lambda phase, done, total: print(f"{phase}: {done}/{total}"),
)

with graph.transaction():
    graph.insert_node(a)
    graph.insert_node(b)
    graph.insert_edge(Edge(source_id=a.id, target_id=b.id, relation="x"))
# Exception inside the block → ROLLBACK. Clean exit → COMMIT.

# Pass 1 — during ingest, be permissive to avoid fragmenting entities
extract_and_merge_parallel(graph, chunks, similarity_threshold=0.88, max_workers=50)
# Pass 2 — after ingest, resolve residual duplicates more strictly
clusters = graph.resolve_duplicates(
    match_keys=["name"],
    match_type=True,
    similarity_threshold=0.9,
    vector_threshold=0.15,
)
for c in clusters:
    print(f"Kept {c.survivor.properties['name']}, merged {len(c.merged)} dups")

prev = None
for chunk in chunks:
    node = Node(type="Chunk", properties={"text": chunk}, vector=embed(chunk))
    graph.insert_node(node)
    if prev is not None:
        graph.insert_edge(Edge(source_id=prev.id, target_id=node.id, relation="followed_by"))
    prev = node

result = graph.retrieve(query=q, query_vector=qv, max_hops=2, max_tokens=4000)
print(result.context_text)   # See exactly what the LLM would receive
# Tune max_hops / max_tokens / search_limit before wiring ask()

python3 -m pytest tests/tests.py -v