Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -45,4 +45,4 @@ node_modules/
/benchmarks/public_comparison/
/scripts/bench/sweep_*.json
benchmarks/baselines/locomo_baseline.json
data/bm25_index.db
data/
2 changes: 2 additions & 0 deletions app.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ def _parse_viewer_allowed_origins() -> Any:
CONSOLIDATION_FORGET_INTERVAL_SECONDS,
CONSOLIDATION_GRACE_PERIOD_DAYS,
CONSOLIDATION_HISTORY_LIMIT,
CONSOLIDATION_IDENTITY_INTERVAL_SECONDS,
CONSOLIDATION_IMPORTANCE_FLOOR_FACTOR,
CONSOLIDATION_IMPORTANCE_PROTECTION_THRESHOLD,
CONSOLIDATION_MIN_CLUSTER_SIZE,
Expand Down Expand Up @@ -363,6 +364,7 @@ def require_api_token() -> None:
creative_interval_seconds=CONSOLIDATION_CREATIVE_INTERVAL_SECONDS,
cluster_interval_seconds=CONSOLIDATION_CLUSTER_INTERVAL_SECONDS,
forget_interval_seconds=CONSOLIDATION_FORGET_INTERVAL_SECONDS,
identity_interval_seconds=CONSOLIDATION_IDENTITY_INTERVAL_SECONDS,
Comment thread
coderabbitai[bot] marked this conversation as resolved.
delete_threshold=CONSOLIDATION_DELETE_THRESHOLD,
archive_threshold=CONSOLIDATION_ARCHIVE_THRESHOLD,
grace_period_days=CONSOLIDATION_GRACE_PERIOD_DAYS,
Expand Down
248 changes: 248 additions & 0 deletions automem/api/entity.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,248 @@
"""Entity API endpoints for AutoMem.

Provides CRUD and query operations for first-class Entity nodes.
"""

from __future__ import annotations

import json
import logging
from typing import Any, Callable, Dict, List, Optional

from flask import Blueprint, abort, jsonify, request

logger = logging.getLogger(__name__)


def _parse_aliases(raw: Any) -> List[str]:
"""Safely parse aliases from a graph result (may be list or JSON string)."""
if not raw:
return []
if isinstance(raw, list):
return raw
if isinstance(raw, str):
try:
parsed = json.loads(raw)
return parsed if isinstance(parsed, list) else []
except Exception:
return []
return []


def _serialize_entity_row(row: list) -> Dict[str, Any]:
"""Convert a result row into a serialisable entity dict.

Expected column order (12 columns):
0: e.id, 1: e.slug, 2: e.category, 3: e.name, 4: e.aliases,
5: e.identity, 6: e.identity_version, 7: e.identity_updated_at,
8: e.identity_source_count, 9: ref_count,
10: e.created_at, 11: e.last_referenced_at
"""
return {
"id": row[0],
"slug": row[1],
"category": row[2],
"name": row[3],
"aliases": _parse_aliases(row[4]),
"identity": row[5],
"identity_version": int(row[6] or 0),
"identity_updated_at": row[7],
"identity_source_count": int(row[8] or 0),
"reference_count": int(row[9] or 0),
"created_at": row[10] if len(row) > 10 else None,
"last_referenced_at": row[11] if len(row) > 11 else None,
}


def create_entity_blueprint(
get_memory_graph: Callable[[], Any],
logger: Any,
require_admin_token_fn: Optional[Callable[[], None]] = None,
) -> Blueprint:
"""Create the entity API blueprint.

Args:
get_memory_graph: Factory that returns the FalkorDB graph instance.
logger: Logger instance.
require_admin_token_fn: Optional callable that aborts if the request
lacks a valid admin token. Used to gate write operations (merge).
"""
bp = Blueprint("entity", __name__)
Comment thread
jescalan marked this conversation as resolved.

@bp.route("/entities", methods=["GET"])
def list_entities() -> Any:
graph = get_memory_graph()
if graph is None:
abort(503, description="FalkorDB is unavailable")

category = request.args.get("category")
try:
limit = min(int(request.args.get("limit", 100)), 500)
except (ValueError, TypeError):
abort(400, description="Invalid limit parameter — must be an integer")

if category:
result = graph.query(
"""
MATCH (e:Entity)
WHERE e.merged_into IS NULL AND e.category = $category
OPTIONAL MATCH (e)-[ref:REFERENCED_IN]->()
WITH e, count(ref) as ref_count
RETURN e.id, e.slug, e.category, e.name, e.aliases,
e.identity, e.identity_version, e.identity_updated_at,
e.identity_source_count, ref_count,
e.created_at, e.last_referenced_at
ORDER BY ref_count DESC
LIMIT $limit
""",
{"category": category, "limit": limit},
)
else:
result = graph.query(
"""
MATCH (e:Entity)
WHERE e.merged_into IS NULL
OPTIONAL MATCH (e)-[ref:REFERENCED_IN]->()
WITH e, count(ref) as ref_count
RETURN e.id, e.slug, e.category, e.name, e.aliases,
e.identity, e.identity_version, e.identity_updated_at,
e.identity_source_count, ref_count,
e.created_at, e.last_referenced_at
ORDER BY ref_count DESC
LIMIT $limit
""",
{"limit": limit},
)

entities = []
for row in getattr(result, "result_set", []) or []:
entities.append(_serialize_entity_row(row))

return jsonify({"status": "success", "entities": entities, "count": len(entities)})

@bp.route("/entity/<slug>", methods=["GET"])
def get_entity(slug: str) -> Any:
graph = get_memory_graph()
if graph is None:
abort(503, description="FalkorDB is unavailable")

# Try direct slug match first, then alias lookup
result = graph.query(
"""
MATCH (e:Entity)
WHERE e.merged_into IS NULL AND (e.slug = $slug OR $slug IN e.aliases)
OPTIONAL MATCH (e)-[ref:REFERENCED_IN]->()
WITH e, count(ref) as ref_count
RETURN e.id, e.slug, e.category, e.name, e.aliases,
e.identity, e.identity_version, e.identity_updated_at,
e.identity_source_count, ref_count,
e.created_at, e.last_referenced_at
LIMIT 1
""",
{"slug": slug},
)

rows = getattr(result, "result_set", []) or []
if not rows:
abort(404, description=f"Entity '{slug}' not found")

entity = _serialize_entity_row(rows[0])
return jsonify({"status": "success", "entity": entity})

@bp.route("/entities/merge-candidates", methods=["GET"])
def merge_candidates() -> Any:
graph = get_memory_graph()
if graph is None:
abort(503, description="FalkorDB is unavailable")

try:
from automem.consolidation.entity_dedup import find_merge_candidates

auto_merge, review = find_merge_candidates(graph)
auto_merge_ids = {(c.entity_a_id, c.entity_b_id) for c in auto_merge}
candidates = []
for c in auto_merge + review:
candidates.append(
{
"entity_a": c.entity_a_id,
"entity_b": c.entity_b_id,
"canonical": c.canonical_id,
"alias": c.alias_id,
"confidence": c.confidence,
"reason": c.reason,
"auto_merge": (c.entity_a_id, c.entity_b_id) in auto_merge_ids,
}
)
return jsonify(
{
"status": "success",
"candidates": candidates,
"count": len(candidates),
}
)
except Exception as exc:
logger.exception("Failed to find merge candidates")
return (
jsonify({"error": "Failed to find merge candidates", "details": str(exc)}),
500,
)

@bp.route("/entity/<slug>/merge", methods=["POST"])
def merge_entity(slug: str) -> Any:
"""Merge one entity into another (admin-only)."""
if require_admin_token_fn is not None:
require_admin_token_fn()

graph = get_memory_graph()
if graph is None:
abort(503, description="FalkorDB is unavailable")

data = request.get_json(silent=True) or {}
target_slug = data.get("merge_into")
if not target_slug:
abort(400, description="'merge_into' slug is required")

# Resolve both entities
alias_result = graph.query(
"MATCH (e:Entity) WHERE e.slug = $slug AND e.merged_into IS NULL RETURN e.id LIMIT 1",
{"slug": slug},
)
canonical_result = graph.query(
"MATCH (e:Entity) WHERE e.slug = $slug AND e.merged_into IS NULL RETURN e.id LIMIT 1",
{"slug": target_slug},
)

alias_rows = getattr(alias_result, "result_set", []) or []
canonical_rows = getattr(canonical_result, "result_set", []) or []

if not alias_rows:
abort(404, description=f"Entity '{slug}' not found")
if not canonical_rows:
abort(404, description=f"Target entity '{target_slug}' not found")

alias_id = alias_rows[0][0]
canonical_id = canonical_rows[0][0]

if alias_id == canonical_id:
abort(400, description="Source and target entity must be different")

try:
from automem.consolidation.entity_dedup import merge_entities

merge_result = merge_entities(graph, canonical_id, alias_id)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
return jsonify(
{
"status": "success",
"merge": {
"canonical": merge_result.canonical_id,
"alias": merge_result.alias_id,
"alias_slug": merge_result.alias_slug,
"edges_moved": merge_result.edges_moved,
},
}
)
except Exception as exc:
logger.exception("Merge failed")
return jsonify({"error": "Merge failed", "details": str(exc)}), 500

return bp
57 changes: 57 additions & 0 deletions automem/api/recall.py
Original file line number Diff line number Diff line change
Expand Up @@ -1947,6 +1947,61 @@ def _run_single_query(
result["jit_enriched"] = True
jit_enriched_count += 1

# Entity identity injection: look up Entity nodes for entities in query/results
entity_identities: List[Dict[str, Any]] = []
if graph is not None and (query_text or results):
try:
query_entities = _extract_query_entities(query_text) if query_text else []
result_entity_slugs: Set[str] = set()
for res in results[:10]:
mem = res.get("memory") or res
for tag in mem.get("tags") or []:
if isinstance(tag, str) and tag.startswith("entity:"):
parts = tag.split(":")
if len(parts) >= 3:
result_entity_slugs.add(parts[2])
all_entity_slugs: Set[str] = set()
for ent_name in query_entities:
all_entity_slugs.add(ent_name.lower().replace(" ", "-"))
all_entity_slugs.update(result_entity_slugs)

slug_list = list(all_entity_slugs)[:10]
try:
ent_result = graph.query(
"""
MATCH (e:Entity)
WHERE e.merged_into IS NULL
AND e.identity IS NOT NULL
AND (e.slug IN $slugs OR any(a IN e.aliases WHERE a IN $slugs))
RETURN e.id, e.slug, e.category, e.name, e.aliases,
e.identity, e.identity_source_count,
e.identity_updated_at
""",
{"slugs": slug_list},
)
for row in getattr(ent_result, "result_set", []) or []:
aliases = row[4] or []
if isinstance(aliases, str):
try:
aliases = json.loads(aliases)
except Exception:
aliases = []
entity_identities.append(
{
"id": row[0],
"name": row[3],
"category": row[2],
"aliases": aliases,
"identity": row[5],
"identity_source_count": int(row[6] or 0),
"identity_updated_at": row[7],
}
)
except Exception:
logger.debug("Entity identity batch lookup failed for slugs %s", slug_list)
except Exception:
logger.debug("Entity identity injection failed")

response = {
"status": "success",
"query": query_text,
Expand Down Expand Up @@ -2000,6 +2055,8 @@ def _run_single_query(
"filtered_count": pre_filter_count - len(results),
}
response["query_time_ms"] = round((time.perf_counter() - query_start) * 1000, 2)
if entity_identities:
response["entities"] = entity_identities
if any_context_profile:
response["context_priority"] = {
"language": any_context_profile.get("language"),
Expand Down
8 changes: 8 additions & 0 deletions automem/api/runtime_bootstrap.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from automem.api.backup import create_backup_blueprint
from automem.api.consolidation import create_consolidation_blueprint_full
from automem.api.enrichment import create_enrichment_blueprint
from automem.api.entity import create_entity_blueprint
from automem.api.graph import create_graph_blueprint
from automem.api.health import create_health_blueprint
from automem.api.memory import create_memory_blueprint_full
Expand Down Expand Up @@ -182,6 +183,12 @@ def register_blueprints(
require_api_token=require_api_token_fn,
)

entity_bp = create_entity_blueprint(
get_memory_graph_fn,
logger,
require_admin_token_fn=require_admin_token_fn,
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.

app.register_blueprint(health_bp)
app.register_blueprint(enrichment_bp)
app.register_blueprint(memory_bp)
Expand All @@ -191,6 +198,7 @@ def register_blueprints(
app.register_blueprint(consolidation_bp)
app.register_blueprint(graph_bp)
app.register_blueprint(stream_bp)
app.register_blueprint(entity_bp)

if is_viewer_enabled():
viewer_bp = create_viewer_blueprint()
Expand Down
6 changes: 6 additions & 0 deletions automem/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,11 +76,17 @@
CONSOLIDATION_CONTROL_LABEL = "ConsolidationControl"
CONSOLIDATION_RUN_LABEL = "ConsolidationRun"
CONSOLIDATION_CONTROL_NODE_ID = os.getenv("CONSOLIDATION_CONTROL_NODE_ID", "global")
CONSOLIDATION_IDENTITY_INTERVAL_SECONDS = int(
os.getenv("CONSOLIDATION_IDENTITY_INTERVAL_SECONDS", str(604800)) # 7 days
)
IDENTITY_SYNTHESIS_MODEL = os.getenv("IDENTITY_SYNTHESIS_MODEL", "gpt-5.4")

CONSOLIDATION_TASK_FIELDS = {
"decay": "decay_last_run",
"creative": "creative_last_run",
"cluster": "cluster_last_run",
"forget": "forget_last_run",
"identity": "identity_last_run",
"full": "full_last_run",
}

Expand Down
Loading
Loading