Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
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
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,4 @@ node_modules/
/benchmarks/snapshots/
/benchmarks/results/
benchmarks/baselines/locomo_baseline.json
data/
2 changes: 2 additions & 0 deletions app.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ def _parse_viewer_allowed_origins() -> Any:
CONSOLIDATION_DELETE_THRESHOLD,
CONSOLIDATION_FORGET_INTERVAL_SECONDS,
CONSOLIDATION_GRACE_PERIOD_DAYS,
CONSOLIDATION_IDENTITY_INTERVAL_SECONDS,
CONSOLIDATION_HISTORY_LIMIT,
CONSOLIDATION_IMPORTANCE_FLOOR_FACTOR,
CONSOLIDATION_IMPORTANCE_PROTECTION_THRESHOLD,
Expand Down Expand Up @@ -358,6 +359,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
218 changes: 218 additions & 0 deletions automem/api/entity.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,218 @@
"""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,
) -> Blueprint:
"""Create the entity API blueprint."""

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")
limit = min(int(request.args.get("limit", 100)), 500)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

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:
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]

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
55 changes: 55 additions & 0 deletions automem/api/recall.py
Original file line number Diff line number Diff line change
Expand Up @@ -1507,6 +1507,59 @@ 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 query_text and graph is not None:
try:
query_entities = _extract_query_entities(query_text)
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)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

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 @@ -1556,6 +1609,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
7 changes: 7 additions & 0 deletions automem/api/runtime_bootstrap.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from automem.api.admin import create_admin_blueprint_full
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 @@ -172,6 +173,11 @@ def register_blueprints(
require_api_token=require_api_token_fn,
)

entity_bp = create_entity_blueprint(
get_memory_graph_fn,
logger,
)
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 @@ -180,6 +186,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 @@ -68,11 +68,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