diff --git a/.gitignore b/.gitignore index 7b9c2006..25fb44f0 100644 --- a/.gitignore +++ b/.gitignore @@ -45,4 +45,4 @@ node_modules/ /benchmarks/public_comparison/ /scripts/bench/sweep_*.json benchmarks/baselines/locomo_baseline.json -data/bm25_index.db +data/ diff --git a/app.py b/app.py index 4ba2be5d..776ed07f 100644 --- a/app.py +++ b/app.py @@ -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, @@ -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, delete_threshold=CONSOLIDATION_DELETE_THRESHOLD, archive_threshold=CONSOLIDATION_ARCHIVE_THRESHOLD, grace_period_days=CONSOLIDATION_GRACE_PERIOD_DAYS, diff --git a/automem/api/entity.py b/automem/api/entity.py new file mode 100644 index 00000000..50b14ae3 --- /dev/null +++ b/automem/api/entity.py @@ -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__) + + @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/", 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//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) + 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 diff --git a/automem/api/recall.py b/automem/api/recall.py index 9e568c45..1766b4aa 100644 --- a/automem/api/recall.py +++ b/automem/api/recall.py @@ -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, @@ -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"), diff --git a/automem/api/runtime_bootstrap.py b/automem/api/runtime_bootstrap.py index 88d929ee..a15da06b 100644 --- a/automem/api/runtime_bootstrap.py +++ b/automem/api/runtime_bootstrap.py @@ -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 @@ -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, + ) + app.register_blueprint(health_bp) app.register_blueprint(enrichment_bp) app.register_blueprint(memory_bp) @@ -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() diff --git a/automem/config.py b/automem/config.py index 6ceaa6c9..715bfa30 100644 --- a/automem/config.py +++ b/automem/config.py @@ -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", } diff --git a/automem/consolidation/entity_dedup.py b/automem/consolidation/entity_dedup.py new file mode 100644 index 00000000..ce530284 --- /dev/null +++ b/automem/consolidation/entity_dedup.py @@ -0,0 +1,277 @@ +"""Entity deduplication logic for AutoMem. + +Identifies and merges duplicate entities by analysing: +- String similarity of slugs (substring, Levenshtein distance) +- Memory overlap: if >60% of memories for entity A are also tagged with entity B +""" + +from __future__ import annotations + +import json +import logging +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set, Tuple + +logger = logging.getLogger(__name__) + + +@dataclass +class MergeCandidate: + """A pair of entities that may be duplicates.""" + + entity_a_id: str + entity_b_id: str + canonical_id: str + alias_id: str + confidence: float + reason: str + + +@dataclass +class MergeResult: + """Outcome of a merge operation.""" + + canonical_id: str + alias_id: str + alias_slug: str + edges_moved: int = 0 + + +def _levenshtein(s1: str, s2: str) -> int: + """Compute Levenshtein edit distance between two strings.""" + if len(s1) < len(s2): + return _levenshtein(s2, s1) + if len(s2) == 0: + return len(s1) + prev_row = list(range(len(s2) + 1)) + for i, c1 in enumerate(s1): + curr_row = [i + 1] + for j, c2 in enumerate(s2): + cost = 0 if c1 == c2 else 1 + curr_row.append(min(curr_row[j] + 1, prev_row[j + 1] + 1, prev_row[j] + cost)) + prev_row = curr_row + return prev_row[-1] + + +def _slug_similarity(slug_a: str, slug_b: str) -> float: + """Heuristic similarity score (0-1) between two entity slugs.""" + if slug_a == slug_b: + return 1.0 + # Substring match — strong signal, especially for hyphenated slugs + # e.g. "alice" in "alice-smith" should score high + if slug_a in slug_b or slug_b in slug_a: + shorter = min(len(slug_a), len(slug_b)) + longer = max(len(slug_a), len(slug_b)) + # Boost substring matches: at least 0.6 similarity + ratio = shorter / longer if longer > 0 else 0.0 + return max(0.6, ratio) + # Levenshtein + max_len = max(len(slug_a), len(slug_b)) + if max_len == 0: + return 1.0 + dist = _levenshtein(slug_a, slug_b) + return max(0.0, 1.0 - dist / max_len) + + +def _memory_overlap(memories_a: Set[str], memories_b: Set[str]) -> float: + """Fraction of memories in the *smaller* set that also appear in the larger set.""" + if not memories_a or not memories_b: + return 0.0 + smaller, larger = ( + (memories_a, memories_b) if len(memories_a) <= len(memories_b) else (memories_b, memories_a) + ) + overlap = len(smaller & larger) + return overlap / len(smaller) + + +def find_merge_candidates( + graph: Any, + *, + min_slug_similarity: float = 0.5, + min_overlap_for_auto: float = 0.6, +) -> Tuple[List[MergeCandidate], List[MergeCandidate]]: + """Scan Entity nodes and find dedup candidates. + + Returns: + (auto_merge, review) - two lists of MergeCandidate objects. + auto_merge: high-confidence candidates that should be merged automatically. + review: lower-confidence candidates for human review. + """ + # Load all entities + result = graph.query( + "MATCH (e:Entity) WHERE e.merged_into IS NULL RETURN e.id, e.slug, e.category, e.aliases" + ) + entities: List[Dict[str, Any]] = [] + for row in getattr(result, "result_set", []) or []: + aliases = row[3] if row[3] else [] + if isinstance(aliases, str): + try: + aliases = json.loads(aliases) + except Exception: + aliases = [] + entities.append( + { + "id": row[0], + "slug": row[1], + "category": row[2], + "aliases": aliases, + } + ) + + if len(entities) < 2: + return [], [] + + # Gather memory sets per entity (single batch query) + entity_memories: Dict[str, Set[str]] = defaultdict(set) + all_edges = graph.query( + "MATCH (e:Entity)-[:REFERENCED_IN]->(m:Memory) WHERE e.merged_into IS NULL RETURN e.id, m.id" + ) + for row in getattr(all_edges, "result_set", []) or []: + if row[0] and row[1]: + entity_memories[row[0]].add(str(row[1])) + + auto_merge: List[MergeCandidate] = [] + review: List[MergeCandidate] = [] + + # Compare pairs within the same category + for i, ea in enumerate(entities): + for eb in entities[i + 1 :]: + if ea["category"] != eb["category"]: + continue + + slug_sim = _slug_similarity(ea["slug"], eb["slug"]) + if slug_sim < min_slug_similarity: + continue + + overlap = _memory_overlap( + entity_memories.get(ea["id"], set()), + entity_memories.get(eb["id"], set()), + ) + + # Determine canonical: longer slug wins + if len(ea["slug"]) >= len(eb["slug"]): + canonical_id, alias_id = ea["id"], eb["id"] + else: + canonical_id, alias_id = eb["id"], ea["id"] + + is_substring = ea["slug"] in eb["slug"] or eb["slug"] in ea["slug"] + confidence = min(1.0, slug_sim * 0.4 + overlap * 0.6) + + candidate = MergeCandidate( + entity_a_id=ea["id"], + entity_b_id=eb["id"], + canonical_id=canonical_id, + alias_id=alias_id, + confidence=confidence, + reason=f"slug_sim={slug_sim:.2f}, overlap={overlap:.2f}, substring={is_substring}", + ) + + if is_substring and overlap > min_overlap_for_auto: + auto_merge.append(candidate) + elif confidence >= 0.5: + review.append(candidate) + + return auto_merge, review + + +def merge_entities(graph: Any, canonical_id: str, alias_id: str) -> MergeResult: + """Merge alias entity into canonical entity. + + - Moves REFERENCED_IN edges from alias to canonical + - Adds alias slug to canonical's aliases list + - Marks alias entity with merged_into + """ + # Get alias slug + res = graph.query( + "MATCH (e:Entity {id: $id}) RETURN e.slug, e.aliases", + {"id": alias_id}, + ) + alias_slug = "" + alias_aliases: List[str] = [] + for row in getattr(res, "result_set", []) or []: + alias_slug = row[0] or "" + raw_aliases = row[1] or [] + if isinstance(raw_aliases, str): + try: + raw_aliases = json.loads(raw_aliases) + except Exception: + raw_aliases = [] + alias_aliases = raw_aliases + + # Move edges: collect memory IDs from alias and batch-create edges on canonical + edge_res = graph.query( + "MATCH (e:Entity {id: $id})-[:REFERENCED_IN]->(m:Memory) RETURN m.id", + {"id": alias_id}, + ) + mem_ids: List[str] = [] + for row in getattr(edge_res, "result_set", []) or []: + if row[0]: + mem_ids.append(str(row[0])) + + edges_moved = len(mem_ids) + if mem_ids: + # Copy edges to canonical + graph.query( + """ + MATCH (e:Entity {id: $canonical_id}) + UNWIND $mem_ids AS mid + MATCH (m:Memory {id: mid}) + MERGE (e)-[:REFERENCED_IN]->(m) + """, + {"canonical_id": canonical_id, "mem_ids": mem_ids}, + ) + # Remove old edges from alias + graph.query( + "MATCH (e:Entity {id: $alias_id})-[r:REFERENCED_IN]->() DELETE r", + {"alias_id": alias_id}, + ) + + # Update canonical aliases (deduplicated) + # Fetch current canonical aliases first + canon_res = graph.query( + "MATCH (e:Entity {id: $id}) RETURN e.aliases", + {"id": canonical_id}, + ) + current_aliases: List[str] = [] + for crow in getattr(canon_res, "result_set", []) or []: + raw = crow[0] or [] + if isinstance(raw, str): + try: + raw = json.loads(raw) + except Exception: + raw = [] + current_aliases = raw + + incoming_aliases = [alias_slug] + alias_aliases if alias_slug else alias_aliases + merged_aliases = list( + dict.fromkeys(current_aliases + incoming_aliases) + ) # dedup, preserve order + graph.query( + """ + MATCH (e:Entity {id: $id}) + SET e.aliases = $aliases + """, + {"id": canonical_id, "aliases": merged_aliases}, + ) + + # Mark alias as merged + graph.query( + "MATCH (e:Entity {id: $id}) SET e.merged_into = $canonical_id", + {"id": alias_id, "canonical_id": canonical_id}, + ) + + logger.info( + "Merged entity %s into %s (slug=%s, edges=%d)", + alias_id, + canonical_id, + alias_slug, + edges_moved, + ) + + return MergeResult( + canonical_id=canonical_id, + alias_id=alias_id, + alias_slug=alias_slug, + edges_moved=edges_moved, + ) diff --git a/automem/consolidation/identity_synthesis.py b/automem/consolidation/identity_synthesis.py new file mode 100644 index 00000000..632df682 --- /dev/null +++ b/automem/consolidation/identity_synthesis.py @@ -0,0 +1,332 @@ +"""Identity synthesis for Entity nodes. + +Gathers memories linked to an entity and uses an LLM to produce a stable +2-5 sentence identity definition. Stores the result on the Entity node. +""" + +from __future__ import annotations + +import logging +from datetime import datetime, timezone +from typing import Any, Callable, Dict, List, Optional + +from automem.config import IDENTITY_SYNTHESIS_MODEL + +logger = logging.getLogger(__name__) + +IDENTITY_PROMPT_TEMPLATE = """\ +Given these memories about {entity_name} ({category}), synthesize a stable identity definition. + +Focus on: +- What/who this is (core definition) +- Key relationships to other entities +- Enduring attributes (not recent events) +- Role/function in the user's life + +Do NOT include: +- Specific dated events (those are episodic) +- Recent/temporary states +- Speculation beyond what the memories support + +Memories: +{memories} + +{previous_section} +Identity (2-5 sentences):""" + + +def _build_previous_section(previous_identity: Optional[str], version: int) -> str: + """Build the previous identity section for the prompt. + + Every 5th version (v5, v10, ...) triggers a full re-synthesis from scratch + to prevent identity drift. The check uses the *current* version before + increment, so v5 triggers re-synthesis and the result is stored as v6. + """ + if not previous_identity: + return "" + if version > 0 and version % 5 == 0: + return "(Full re-synthesis — ignore previous identity and synthesize fresh from memories.)" + return f"Previous identity (refine, don't replace unless contradicted):\n{previous_identity}" + + +def _gather_entity_memories( + graph: Any, + entity_id: str, + limit: int = 50, +) -> List[Dict[str, Any]]: + """Fetch memories linked to an entity, ordered by importance desc.""" + result = graph.query( + """ + MATCH (e:Entity {id: $id})-[:REFERENCED_IN]->(m:Memory) + RETURN m.id, m.content, m.importance, m.timestamp, m.type + ORDER BY coalesce(m.importance, 0.0) DESC + LIMIT $limit + """, + {"id": entity_id, "limit": limit}, + ) + memories: List[Dict[str, Any]] = [] + for row in getattr(result, "result_set", []) or []: + memories.append( + { + "id": row[0], + "content": row[1] or "", + "importance": row[2], + "timestamp": row[3], + "type": row[4], + } + ) + return memories + + +def _format_memories_for_prompt(memories: List[Dict[str, Any]]) -> str: + """Format memories into a text block for the LLM prompt.""" + lines: List[str] = [] + for i, mem in enumerate(memories, 1): + content = (mem.get("content") or "").strip() + if not content: + continue + ts = mem.get("timestamp") or "unknown" + lines.append(f"{i}. [{ts[:10]}] {content[:300]}") + return "\n".join(lines) + + +def synthesize_identity( + graph: Any, + entity_id: str, + openai_client: Any, + *, + model: Optional[str] = None, + memory_limit: int = 50, +) -> Optional[str]: + """Synthesize identity for a single entity. + + Args: + graph: FalkorDB graph instance. + entity_id: The entity ID (e.g. "entity:people:alice-smith"). + openai_client: OpenAI-compatible client. + model: Model override (defaults to IDENTITY_SYNTHESIS_MODEL config). + memory_limit: Max memories to gather. + + Returns: + The synthesized identity string, or None if insufficient data. + """ + model = model or IDENTITY_SYNTHESIS_MODEL + + # Fetch entity details + ent_result = graph.query( + """ + MATCH (e:Entity {id: $id}) + RETURN e.name, e.category, e.identity, e.identity_version + """, + {"id": entity_id}, + ) + ent_rows = getattr(ent_result, "result_set", []) or [] + if not ent_rows: + logger.warning("Entity %s not found", entity_id) + return None + + entity_name = ent_rows[0][0] or entity_id + category = ent_rows[0][1] or "unknown" + previous_identity = ent_rows[0][2] + current_version = int(ent_rows[0][3] or 0) + + # Gather memories + memories = _gather_entity_memories(graph, entity_id, limit=memory_limit) + if not memories: + logger.info("No memories for entity %s, skipping synthesis", entity_id) + return None + + formatted_memories = _format_memories_for_prompt(memories) + previous_section = _build_previous_section(previous_identity, current_version) + + prompt = IDENTITY_PROMPT_TEMPLATE.format( + entity_name=entity_name, + category=category, + memories=formatted_memories, + previous_section=previous_section, + ) + + try: + # Use max_completion_tokens (newer API) with max_tokens fallback + try: + response = openai_client.chat.completions.create( + model=model, + messages=[ + { + "role": "system", + "content": "You synthesize concise identity definitions from episodic memories.", + }, + {"role": "user", "content": prompt}, + ], + temperature=0.3, + max_completion_tokens=500, + ) + except TypeError: + # Older clients that don't support max_completion_tokens + response = openai_client.chat.completions.create( + model=model, + messages=[ + { + "role": "system", + "content": "You synthesize concise identity definitions from episodic memories.", + }, + {"role": "user", "content": prompt}, + ], + temperature=0.3, + max_tokens=500, + ) + raw_content = response.choices[0].message.content or "" + identity_text = raw_content.strip() + if not identity_text: + logger.warning("LLM returned empty content for entity %s", entity_id) + return None + except Exception: + logger.exception("LLM call failed for entity %s", entity_id) + return None + + # Get the true total reference count (not truncated by memory_limit) + ref_count_result = graph.query( + "MATCH (e:Entity {id: $id})-[:REFERENCED_IN]->(m:Memory) RETURN count(m)", + {"id": entity_id}, + ) + ref_rows = getattr(ref_count_result, "result_set", []) or [] + try: + total_ref_count = int(ref_rows[0][0]) if ref_rows else len(memories) + except (ValueError, TypeError, IndexError): + total_ref_count = len(memories) + + # Store on entity node + now = datetime.now(timezone.utc).isoformat() + graph.query( + """ + MATCH (e:Entity {id: $id}) + SET e.identity = $identity, + e.identity_version = $version, + e.identity_updated_at = $now, + e.identity_source_count = $source_count + """, + { + "id": entity_id, + "identity": identity_text, + "version": current_version + 1, + "now": now, + "source_count": int(total_ref_count), + }, + ) + + logger.info( + "Synthesized identity for %s (v%d, %d sources, %d chars)", + entity_id, + current_version + 1, + len(memories), + len(identity_text), + ) + return identity_text + + +def run_identity_consolidation( + graph: Any, + openai_client: Any, + *, + model: Optional[str] = None, + min_references: int = 1, + dry_run: bool = False, +) -> Dict[str, Any]: + """Run identity synthesis for all eligible entities. + + This includes: + 1. Entity dedup (find and auto-merge candidates) + 2. Identity synthesis for entities needing it + + Returns: + Summary dict with counts and details. + """ + from automem.consolidation.entity_dedup import find_merge_candidates, merge_entities + + result: Dict[str, Any] = { + "merges_performed": [], + "merge_candidates_for_review": [], + "identities_synthesized": 0, + "entities_examined": 0, + "errors": [], + } + + # Step 1: Dedup + try: + auto_merge, review = find_merge_candidates(graph) + result["merge_candidates_for_review"] = [ + { + "canonical": c.canonical_id, + "alias": c.alias_id, + "confidence": c.confidence, + "reason": c.reason, + } + for c in review + ] + + if not dry_run: + for candidate in auto_merge: + try: + merge_result = merge_entities(graph, candidate.canonical_id, candidate.alias_id) + result["merges_performed"].append( + { + "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( + "Failed to merge %s into %s", + candidate.alias_id, + candidate.canonical_id, + ) + result["errors"].append(f"merge {candidate.alias_id}: {exc}") + except Exception as exc: + logger.exception("Entity dedup failed") + result["errors"].append(f"dedup: {exc}") + + # Step 2: Synthesize identities (only for new or changed entities) + ent_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.identity_source_count, e.identity, ref_count + """ + ) + for row in getattr(ent_result, "result_set", []) or []: + entity_id = row[0] + stored_source_count = int(row[1] or 0) + existing_identity = row[2] + actual_ref_count = int(row[3] or 0) + result["entities_examined"] += 1 + + if actual_ref_count < min_references: + continue + + # Only re-synthesize if entity has no identity or reference count changed + needs_synthesis = existing_identity is None or actual_ref_count != stored_source_count + if not needs_synthesis: + continue + + if dry_run: + result["identities_synthesized"] += 1 + continue + + try: + identity = synthesize_identity( + graph, + entity_id, + openai_client, + model=model, + ) + if identity: + result["identities_synthesized"] += 1 + except Exception as exc: + logger.exception("Identity synthesis failed for %s", entity_id) + result["errors"].append(f"synthesis {entity_id}: {exc}") + + return result diff --git a/automem/consolidation/runtime_bindings.py b/automem/consolidation/runtime_bindings.py index 1db963b8..d8621830 100644 --- a/automem/consolidation/runtime_bindings.py +++ b/automem/consolidation/runtime_bindings.py @@ -61,6 +61,7 @@ def create_consolidation_runtime( creative_interval_seconds: int, cluster_interval_seconds: int, forget_interval_seconds: int, + identity_interval_seconds: int = 604800, delete_threshold: float, archive_threshold: float, grace_period_days: int, @@ -97,6 +98,7 @@ def _apply_scheduler_overrides(scheduler: ConsolidationScheduler) -> None: creative_interval_seconds=creative_interval_seconds, cluster_interval_seconds=cluster_interval_seconds, forget_interval_seconds=forget_interval_seconds, + identity_interval_seconds=identity_interval_seconds, ) def persist_consolidation_run(graph: Any, result: Dict[str, Any]) -> None: diff --git a/automem/consolidation/runtime_helpers.py b/automem/consolidation/runtime_helpers.py index ffd46dc2..094189bb 100644 --- a/automem/consolidation/runtime_helpers.py +++ b/automem/consolidation/runtime_helpers.py @@ -82,6 +82,7 @@ def apply_scheduler_overrides( creative_interval_seconds: int, cluster_interval_seconds: int, forget_interval_seconds: int, + identity_interval_seconds: int = 604800, ) -> None: """Override default scheduler intervals using configuration.""" overrides = { @@ -89,6 +90,7 @@ def apply_scheduler_overrides( "creative": timedelta(seconds=creative_interval_seconds), "cluster": timedelta(seconds=cluster_interval_seconds), "forget": timedelta(seconds=forget_interval_seconds), + "identity": timedelta(seconds=identity_interval_seconds), } for task, interval in overrides.items(): @@ -99,7 +101,7 @@ def apply_scheduler_overrides( def tasks_for_mode(mode: str, task_fields: Dict[str, str]) -> List[str]: """Map a consolidation mode to its task identifiers.""" if mode == "full": - return ["decay", "creative", "cluster", "forget", "full"] + return ["decay", "creative", "cluster", "forget", "identity", "full"] if mode in task_fields: return [mode] return [mode] diff --git a/automem/consolidation/runtime_scheduler.py b/automem/consolidation/runtime_scheduler.py index ea6669d4..107d0d3d 100644 --- a/automem/consolidation/runtime_scheduler.py +++ b/automem/consolidation/runtime_scheduler.py @@ -41,6 +41,9 @@ def run_consolidation_tick( if "forget" in steps: affected_count += steps["forget"].get("archived", 0) affected_count += steps["forget"].get("deleted", 0) + if "identity" in steps: + affected_count += steps["identity"].get("identities_synthesized", 0) + affected_count += len(steps["identity"].get("merges_performed", [])) elapsed_ms = int((perf_counter_fn() - task_start) * 1000) next_runs = scheduler.get_next_runs() diff --git a/consolidation.py b/consolidation.py index be81815c..f96ccc59 100644 --- a/consolidation.py +++ b/consolidation.py @@ -12,6 +12,7 @@ import json import logging import math +import os import time from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -167,6 +168,21 @@ def __init__( frozenset(protected_types) if protected_types is not None else self.PROTECTED_TYPES ) + def _get_openai_client(self) -> Optional[Any]: + """Obtain an OpenAI-compatible client for LLM calls.""" + try: + import openai + + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + logger.warning("OPENAI_API_KEY not set; identity synthesis unavailable") + return None + base_url = os.environ.get("OPENAI_BASE_URL") + return openai.OpenAI(api_key=api_key, **({"base_url": base_url} if base_url else {})) + except Exception: + logger.exception("Failed to create OpenAI client") + return None + def _query_graph( self, query: str, params: Optional[Dict[str, Any]] = None ) -> List[Sequence[Any]]: @@ -273,7 +289,12 @@ def _should_protect_memory( """Return whether a memory should be protected from archiving/deletion.""" protected_flag = memory.get("protected") if isinstance(protected_flag, str): - protected_flag = protected_flag.strip().lower() not in {"", "0", "false", "no"} + protected_flag = protected_flag.strip().lower() not in { + "", + "0", + "false", + "no", + } if protected_flag: reason = memory.get("protected_reason") or "explicitly protected" return True, str(reason) @@ -299,7 +320,10 @@ def _should_protect_memory( if created_at: age_days = (current_time - created_at).days if age_days < self.grace_period_days: - return True, f"within grace period ({age_days} < {self.grace_period_days} days)" + return ( + True, + f"within grace period ({age_days} < {self.grace_period_days} days)", + ) memory_type = memory.get("type") if memory_type and str(memory_type) in self.protected_types: @@ -395,7 +419,10 @@ def discover_creative_associations(self, sample_size: int = 20) -> List[Dict[str if similarity < 0.3: connection_type = "CONTRASTS_WITH" confidence = 0.6 - elif {mem1.type, mem2.type} == {"Insight", "Pattern"} and similarity > 0.5: + elif {mem1.type, mem2.type} == { + "Insight", + "Pattern", + } and similarity > 0.5: connection_type = "DISCOVERED" connection_kind = "explains" confidence = 0.7 @@ -595,7 +622,13 @@ def apply_controlled_forgetting(self, dry_run: bool = True) -> Dict[str, Any]: Returns statistics about what was archived/deleted. """ - stats = {"examined": 0, "archived": [], "deleted": [], "preserved": 0, "protected": []} + stats = { + "examined": 0, + "archived": [], + "deleted": [], + "preserved": 0, + "protected": [], + } # Get all memories with scores all_memories_query = """ @@ -785,7 +818,10 @@ def apply_controlled_forgetting(self, dry_run: bool = True) -> Dict[str, Any]: return stats def consolidate( - self, mode: str = "full", dry_run: bool = True, decay_threshold: Optional[float] = None + self, + mode: str = "full", + dry_run: bool = True, + decay_threshold: Optional[float] = None, ) -> Dict[str, Any]: """ Run full consolidation cycle. @@ -922,7 +958,11 @@ def consolidate( """ try: self.graph.query( - link_query, {"meta_id": cluster["cluster_id"], "mem_id": mem_id} + link_query, + { + "meta_id": cluster["cluster_id"], + "mem_id": mem_id, + }, ) except: pass @@ -935,7 +975,31 @@ def consolidate( "sample_clusters": clusters[:2] if clusters else [], } - # Step 4: Controlled forgetting + # Step 4: Identity consolidation (entity dedup + identity synthesis) + if mode in ["full", "identity"]: + logger.info("Running identity consolidation...") + try: + from automem.consolidation.identity_synthesis import run_identity_consolidation + + # Get OpenAI client from environment + openai_client = self._get_openai_client() + if openai_client is not None: + identity_result = run_identity_consolidation( + self.graph, + openai_client, + dry_run=dry_run, + ) + results["steps"]["identity"] = identity_result + else: + results["steps"]["identity"] = { + "skipped": True, + "reason": "No OpenAI client available", + } + except Exception as exc: + logger.exception("Identity consolidation failed") + results["steps"]["identity"] = {"error": str(exc)} + + # Step 5: Controlled forgetting if mode in ["full", "forget"]: logger.info("Applying controlled forgetting...") forget_stats = self.apply_controlled_forgetting(dry_run=dry_run) @@ -1064,6 +1128,7 @@ def __init__(self, consolidator: MemoryConsolidator): "creative": {"interval": timedelta(days=7), "last_run": None}, "cluster": {"interval": timedelta(days=30), "last_run": None}, "forget": {"interval": timedelta(days=90), "last_run": None}, + "identity": {"interval": timedelta(days=7), "last_run": None}, } self.history = [] @@ -1112,8 +1177,17 @@ def run_scheduled_tasks( else: result = self.consolidator.consolidate(mode=task_type, dry_run=False) - # Update schedule - self.schedules[task_type]["last_run"] = datetime.now(timezone.utc) + # Don't advance schedule if the identity task was skipped or failed + identity_not_completed = ( + task_type == "identity" + and isinstance(result, dict) + and ( + result.get("steps", {}).get("identity", {}).get("skipped") + or result.get("steps", {}).get("identity", {}).get("error") + ) + ) + if not identity_not_completed: + self.schedules[task_type]["last_run"] = datetime.now(timezone.utc) # Record in history self.history.append( diff --git a/scripts/migrate_entity_nodes.py b/scripts/migrate_entity_nodes.py new file mode 100644 index 00000000..03a2d7eb --- /dev/null +++ b/scripts/migrate_entity_nodes.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +""" +Migrate entity tags on Memory nodes into first-class Entity nodes in FalkorDB. + +Scans all Memory nodes for `entity:{category}:{slug}` tags, creates Entity nodes, +and links them via REFERENCED_IN relationships. Idempotent (safe to re-run). + +Usage: + python scripts/migrate_entity_nodes.py [--dry-run] +""" + +from __future__ import annotations + +import argparse +import logging +import os +import re +import sys +from collections import defaultdict +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +# Ensure project root is on sys.path +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from dotenv import load_dotenv + +load_dotenv() +load_dotenv(Path.home() / ".config" / "automem" / ".env") + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger(__name__) + +ENTITY_TAG_RE = re.compile(r"^entity:([a-z]+):(.+)$") + + +def slug_to_name(slug: str) -> str: + """Convert a slug like 'alice-smith' to 'Alice Smith'.""" + return slug.replace("-", " ").title() + + +def connect_falkordb() -> Any: + """Connect to FalkorDB and return the graph.""" + from automem.config import FALKORDB_PORT, GRAPH_NAME + + try: + from falkordb import FalkorDB + except ImportError: + logger.error("falkordb package not installed") + sys.exit(1) + + host = os.getenv("FALKORDB_HOST", "localhost") + password = os.getenv("FALKORDB_PASSWORD") + params = {"host": host, "port": FALKORDB_PORT} + if password: + params["password"] = password + params["username"] = "default" + + db = FalkorDB(**params) + return db.select_graph(GRAPH_NAME) + + +def collect_entity_tags(graph) -> dict[str, list[str]]: + """Scan all Memory nodes and collect entity tags mapped to memory IDs. + + Returns: + dict mapping entity tag (e.g. "entity:people:alice-smith") to list of memory IDs + """ + result = graph.query("MATCH (m:Memory) WHERE m.tags IS NOT NULL RETURN m.id, m.tags") + entity_to_memories: dict[str, list[str]] = defaultdict(list) + for row in getattr(result, "result_set", []) or []: + mem_id = row[0] + tags = row[1] + if not isinstance(tags, list): + continue + for tag in tags: + if isinstance(tag, str) and ENTITY_TAG_RE.match(tag): + entity_to_memories[tag].append(mem_id) + return dict(entity_to_memories) + + +def run_migration(graph, entity_to_memories: dict[str, list[str]], *, dry_run: bool) -> None: + """Create Entity nodes and REFERENCED_IN edges.""" + now = datetime.now(timezone.utc).isoformat() + created_entities = 0 + created_edges = 0 + + for entity_tag, memory_ids in sorted(entity_to_memories.items()): + match = ENTITY_TAG_RE.match(entity_tag) + if not match: + continue + category = match.group(1) + slug = match.group(2) + name = slug_to_name(slug) + + logger.info( + "Entity: %s (category=%s, slug=%s, references=%d)%s", + entity_tag, + category, + slug, + len(memory_ids), + " [DRY RUN]" if dry_run else "", + ) + + if not dry_run: + # MERGE Entity node (idempotent) + graph.query( + """ + MERGE (e:Entity {id: $id}) + ON CREATE SET + e.slug = $slug, + e.category = $category, + e.name = $name, + e.aliases = [], + e.identity = null, + e.identity_version = 0, + e.identity_updated_at = null, + e.identity_source_count = $ref_count, + e.created_at = $now, + e.last_referenced_at = $now + ON MATCH SET + e.identity_source_count = $ref_count, + e.last_referenced_at = $now + """, + { + "id": entity_tag, + "slug": slug, + "category": category, + "name": name, + "ref_count": len(memory_ids), + "now": now, + }, + ) + created_entities += 1 + + # Create REFERENCED_IN edges (batched) + graph.query( + """ + MATCH (e:Entity {id: $entity_id}) + UNWIND $mem_ids AS mid + MATCH (m:Memory {id: mid}) + MERGE (e)-[:REFERENCED_IN]->(m) + """, + {"entity_id": entity_tag, "mem_ids": memory_ids}, + ) + created_edges += len(memory_ids) + else: + created_entities += 1 + created_edges += len(memory_ids) + + logger.info( + "Summary: %d entities, %d edges%s", + created_entities, + created_edges, + " (dry run)" if dry_run else " created", + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Migrate entity tags to Entity nodes") + parser.add_argument("--dry-run", action="store_true", help="Show what would be created") + args = parser.parse_args() + + graph = connect_falkordb() + entity_to_memories = collect_entity_tags(graph) + logger.info("Found %d distinct entity tags across memories", len(entity_to_memories)) + + if not entity_to_memories: + logger.info("Nothing to migrate.") + return + + run_migration(graph, entity_to_memories, dry_run=args.dry_run) + + +if __name__ == "__main__": + main() diff --git a/tests/contracts/test_routes_contract.py b/tests/contracts/test_routes_contract.py index a4f70032..1d6443c1 100644 --- a/tests/contracts/test_routes_contract.py +++ b/tests/contracts/test_routes_contract.py @@ -32,6 +32,10 @@ ("GET", "/viewer/"), ("GET", "/stream"), ("GET", "/stream/status"), + ("GET", "/entities"), + ("GET", "/entities/merge-candidates"), + ("GET", "/entity/"), + ("POST", "/entity//merge"), } diff --git a/tests/test_entity_identities.py b/tests/test_entity_identities.py new file mode 100644 index 00000000..1fd209cf --- /dev/null +++ b/tests/test_entity_identities.py @@ -0,0 +1,709 @@ +"""Tests for the Entity Identity Synthesis feature. + +Covers: +- Entity node creation from tags (migration) +- Dedup candidate detection (string similarity, co-occurrence) +- Merge operation +- Identity synthesis (mocked LLM) +- Recall with entity injection +""" + +from __future__ import annotations + +import json +import re +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Any, Dict, List, Optional, Set +from unittest.mock import MagicMock, patch + +import pytest + +from tests.support.fake_graph import FakeGraph, FakeNode, FakeResult + +# --------------------------------------------------------------------------- +# Helpers: Entity-aware FakeGraph extension +# --------------------------------------------------------------------------- + + +class EntityFakeGraph(FakeGraph): + """FakeGraph with Entity node support for testing.""" + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.entities: Dict[str, Dict[str, Any]] = {} + self.entity_edges: List[Dict[str, str]] = [] # {"entity_id": ..., "memory_id": ...} + + def query(self, query: str, params: Dict[str, Any] | None = None, **kwargs: Any) -> FakeResult: + params = params or {} + self.queries.append((query, params)) + + # MERGE Entity node + if "MERGE (e:Entity {id: $id})" in query: + eid = params["id"] + if eid not in self.entities: + self.entities[eid] = { + "id": eid, + "slug": params.get("slug", ""), + "category": params.get("category", ""), + "name": params.get("name", ""), + "aliases": params.get("aliases", []), + "identity": params.get("identity"), + "identity_version": params.get("identity_version", 0), + "identity_updated_at": params.get("identity_updated_at"), + "identity_source_count": params.get("identity_source_count", 0), + "merged_into": None, + "created_at": params.get("created_at"), + "last_referenced_at": params.get("last_referenced_at"), + } + else: + # ON MATCH SET updates + ent = self.entities[eid] + if "identity_source_count" in params: + ent["identity_source_count"] = params["identity_source_count"] + if "last_referenced_at" in params: + ent["last_referenced_at"] = params["last_referenced_at"] + if "identity" in params: + ent["identity"] = params["identity"] + if "identity_version" in params: + ent["identity_version"] = params["identity_version"] + return FakeResult([]) + + # MERGE REFERENCED_IN edge (single or batched via UNWIND) + if "MERGE (e)-[:REFERENCED_IN]->(m)" in query: + entity_id = params.get("entity_id") or params.get("canonical_id") + # Batched: UNWIND $mem_ids + mem_ids = params.get("mem_ids") + if entity_id and mem_ids: + for mid in mem_ids: + edge = {"entity_id": entity_id, "memory_id": mid} + if edge not in self.entity_edges: + self.entity_edges.append(edge) + return FakeResult([]) + mem_id = params.get("mem_id") + if entity_id and mem_id: + # Avoid duplicate edges + edge = {"entity_id": entity_id, "memory_id": mem_id} + if edge not in self.entity_edges: + self.entity_edges.append(edge) + return FakeResult([]) + + # Batch entity-memory edge query for dedup + if ( + "MATCH (e:Entity)-[:REFERENCED_IN]->(m:Memory)" in query + and "RETURN e.id, m.id" in query + ): + rows = [] + for edge in self.entity_edges: + ent = self.entities.get(edge["entity_id"]) + if ent and not ent.get("merged_into"): + rows.append([edge["entity_id"], edge["memory_id"]]) + return FakeResult(rows) + + # Entity query: MATCH (e:Entity) ... RETURN e.id, e.slug, ... + if "MATCH (e:Entity)" in query and "e.merged_into IS NULL" in query: + # Batch slug lookup (recall entity injection) + if "e.slug IN $slugs" in query: + slugs = params.get("slugs", []) + rows = [] + for ent in self.entities.values(): + if ent.get("merged_into"): + continue + if ent.get("identity") is None and "e.identity IS NOT NULL" in query: + continue + matched = ent["slug"] in slugs or any( + a in slugs for a in (ent.get("aliases") or []) + ) + if matched: + rows.append( + [ + ent["id"], + ent["slug"], + ent["category"], + ent["name"], + ent.get("aliases", []), + ent.get("identity"), + ent.get("identity_source_count", 0), + ent.get("identity_updated_at"), + ] + ) + return FakeResult(rows) + + # Single slug lookup + if "e.slug = $slug" in query or "$slug IN e.aliases" in query: + slug = params.get("slug", "") + for ent in self.entities.values(): + if ent.get("merged_into"): + continue + if ent["slug"] == slug or slug in (ent.get("aliases") or []): + # Compute ref_count for this entity + ref_count = sum(1 for e in self.entity_edges if e["entity_id"] == ent["id"]) + return FakeResult( + [ + [ + ent["id"], + ent["slug"], + ent["category"], + ent["name"], + ent.get("aliases", []), + ent.get("identity"), + ent.get("identity_version", 0), + ent.get("identity_updated_at"), + ent.get("identity_source_count", 0), + ref_count, + ent.get("created_at"), + ent.get("last_referenced_at"), + ] + ] + ) + return FakeResult([]) + + if "RETURN e.id, e.slug, e.category, e.aliases" in query: + rows = [] + for ent in self.entities.values(): + if ent.get("merged_into"): + continue + rows.append( + [ + ent["id"], + ent["slug"], + ent["category"], + ent.get("aliases", []), + ] + ) + return FakeResult(rows) + + if "RETURN e.id, e.identity_source_count, e.identity" in query: + rows = [] + for ent in self.entities.values(): + if ent.get("merged_into"): + continue + ref_count = sum(1 for e in self.entity_edges if e["entity_id"] == ent["id"]) + rows.append( + [ + ent["id"], + ent.get("identity_source_count", 0), + ent.get("identity"), + ref_count, + ] + ) + return FakeResult(rows) + + # Generic list (with ref_count) + rows = [] + for ent in self.entities.values(): + if ent.get("merged_into"): + continue + ref_count = sum(1 for e in self.entity_edges if e["entity_id"] == ent["id"]) + rows.append( + [ + ent["id"], + ent["slug"], + ent["category"], + ent["name"], + ent.get("aliases", []), + ent.get("identity"), + ent.get("identity_version", 0), + ent.get("identity_updated_at"), + ent.get("identity_source_count", 0), + ref_count, + ent.get("created_at"), + ent.get("last_referenced_at"), + ] + ) + return FakeResult(rows) + + # Entity by ID + if "MATCH (e:Entity {id: $id})" in query: + eid = params.get("id") + ent = self.entities.get(eid) + if not ent: + return FakeResult([]) + + if "REFERENCED_IN" in query: + # Get linked memories + mem_ids = [e["memory_id"] for e in self.entity_edges if e["entity_id"] == eid] + rows = [] + for mid in mem_ids: + mem = self.memories.get(mid) + if mem: + rows.append( + [ + mem.get("id"), + mem.get("content", ""), + mem.get("importance", 0.5), + mem.get("timestamp"), + mem.get("type"), + ] + ) + return FakeResult(rows) + + if "RETURN e.name, e.category, e.identity, e.identity_version" in query: + return FakeResult( + [ + [ + ent["name"], + ent["category"], + ent.get("identity"), + ent.get("identity_version", 0), + ] + ] + ) + + if "RETURN e.slug, e.aliases" in query: + return FakeResult([[ent["slug"], ent.get("aliases", [])]]) + + if "SET e.identity" in query: + ent["identity"] = params.get("identity") + ent["identity_version"] = params.get("version", 0) + ent["identity_updated_at"] = params.get("now") + ent["identity_source_count"] = params.get("source_count", 0) + return FakeResult([]) + + if "SET e.merged_into" in query: + ent["merged_into"] = params.get("canonical_id") + return FakeResult([]) + + if "SET e.aliases" in query: + # New behavior: full replacement (deduplicated in Python) + if "aliases" in params: + ent["aliases"] = params["aliases"] + else: + new_aliases = params.get("new_aliases", []) + ent["aliases"] = (ent.get("aliases") or []) + new_aliases + return FakeResult([]) + + if "RETURN e.aliases" in query: + return FakeResult([[ent.get("aliases", [])]]) + + return FakeResult([[FakeNode(ent)]]) + + # Memory tags query for migration + if "MATCH (m:Memory)" in query and "RETURN m.id, m.tags" in query: + rows = [] + for mem in self.memories.values(): + rows.append([mem["id"], mem.get("tags", [])]) + return FakeResult(rows) + + return super().query(query, params, **kwargs) + + +# --------------------------------------------------------------------------- +# Test: Entity node creation from tags (migration logic) +# --------------------------------------------------------------------------- + + +class TestEntityMigration: + """Test the migration script logic.""" + + def test_collect_entity_tags(self) -> None: + """Entity tags are correctly grouped by tag.""" + from scripts.migrate_entity_nodes import ENTITY_TAG_RE, collect_entity_tags + + graph = EntityFakeGraph() + graph.memories["m1"] = { + "id": "m1", + "content": "Alice is great", + "tags": ["entity:people:alice-smith", "entity:organizations:acme-corp"], + "importance": 0.8, + } + graph.memories["m2"] = { + "id": "m2", + "content": "Alice likes music", + "tags": ["entity:people:alice-smith", "music"], + "importance": 0.6, + } + + result = collect_entity_tags(graph) + assert "entity:people:alice-smith" in result + assert set(result["entity:people:alice-smith"]) == {"m1", "m2"} + assert "entity:organizations:acme-corp" in result + assert result["entity:organizations:acme-corp"] == ["m1"] + + def test_slug_to_name(self) -> None: + from scripts.migrate_entity_nodes import slug_to_name + + assert slug_to_name("alice-smith") == "Alice Smith" + assert slug_to_name("acme-corp") == "Acme Corp" + assert slug_to_name("cool-startup") == "Cool Startup" + + def test_migration_creates_entities(self) -> None: + """Migration creates Entity nodes and edges.""" + from scripts.migrate_entity_nodes import run_migration + + graph = EntityFakeGraph() + graph.memories["m1"] = { + "id": "m1", + "content": "Alice at Acme Corp", + "tags": ["entity:people:alice-smith"], + "importance": 0.8, + } + entity_map = {"entity:people:alice-smith": ["m1"]} + + run_migration(graph, entity_map, dry_run=False) + + assert "entity:people:alice-smith" in graph.entities + ent = graph.entities["entity:people:alice-smith"] + assert ent["slug"] == "alice-smith" + assert ent["category"] == "people" + assert ent["name"] == "Alice Smith" + assert len(graph.entity_edges) == 1 + + def test_migration_dry_run(self) -> None: + """Dry run doesn't create anything.""" + from scripts.migrate_entity_nodes import run_migration + + graph = EntityFakeGraph() + entity_map = {"entity:people:alice-smith": ["m1"]} + + run_migration(graph, entity_map, dry_run=True) + assert len(graph.entities) == 0 + assert len(graph.entity_edges) == 0 + + +# --------------------------------------------------------------------------- +# Test: Entity dedup +# --------------------------------------------------------------------------- + + +class TestEntityDedup: + """Test entity deduplication logic.""" + + def test_levenshtein_distance(self) -> None: + from automem.consolidation.entity_dedup import _levenshtein + + assert _levenshtein("kitten", "sitting") == 3 + assert _levenshtein("abc", "abc") == 0 + assert _levenshtein("", "abc") == 3 + + def test_slug_similarity_exact(self) -> None: + from automem.consolidation.entity_dedup import _slug_similarity + + assert _slug_similarity("alice", "alice") == 1.0 + + def test_slug_similarity_substring(self) -> None: + from automem.consolidation.entity_dedup import _slug_similarity + + sim = _slug_similarity("alice", "alice-smith") + assert sim > 0.2 # alice is substring + + def test_memory_overlap(self) -> None: + from automem.consolidation.entity_dedup import _memory_overlap + + a = {"m1", "m2", "m3"} + b = {"m1", "m2", "m3", "m4", "m5"} + assert _memory_overlap(a, b) == 1.0 # all of a are in b + + c = {"m1", "m6"} + assert _memory_overlap(c, b) == 0.5 # 1 of 2 + + def test_find_candidates_substring_overlap(self) -> None: + """Substring + high overlap → auto-merge candidate.""" + from automem.consolidation.entity_dedup import find_merge_candidates + + graph = EntityFakeGraph() + # Create two entities with substring relationship and shared memories + graph.entities["entity:people:alice"] = { + "id": "entity:people:alice", + "slug": "alice", + "category": "people", + "aliases": [], + "merged_into": None, + } + graph.entities["entity:people:alice-smith"] = { + "id": "entity:people:alice-smith", + "slug": "alice-smith", + "category": "people", + "aliases": [], + "merged_into": None, + } + + # alice references m1, m2, m3 + for mid in ["m1", "m2", "m3"]: + graph.entity_edges.append({"entity_id": "entity:people:alice", "memory_id": mid}) + graph.memories[mid] = {"id": mid, "content": f"Memory {mid}", "tags": []} + + # alice-smith references m1, m2, m3, m4, m5 + for mid in ["m1", "m2", "m3", "m4", "m5"]: + graph.entity_edges.append({"entity_id": "entity:people:alice-smith", "memory_id": mid}) + if mid not in graph.memories: + graph.memories[mid] = { + "id": mid, + "content": f"Memory {mid}", + "tags": [], + } + + auto_merge, _review = find_merge_candidates(graph) + assert len(auto_merge) >= 1 + # The canonical should be the longer slug + assert auto_merge[0].canonical_id == "entity:people:alice-smith" + assert auto_merge[0].alias_id == "entity:people:alice" + + def test_different_categories_not_merged(self) -> None: + """Entities in different categories are not candidates.""" + from automem.consolidation.entity_dedup import find_merge_candidates + + graph = EntityFakeGraph() + graph.entities["entity:people:python"] = { + "id": "entity:people:python", + "slug": "python", + "category": "people", + "aliases": [], + "merged_into": None, + } + graph.entities["entity:tools:python"] = { + "id": "entity:tools:python", + "slug": "python", + "category": "tools", + "aliases": [], + "merged_into": None, + } + auto_merge, _review = find_merge_candidates(graph) + assert len(auto_merge) == 0 + assert len(_review) == 0 + + def test_merge_entities(self) -> None: + """Merge moves edges and updates aliases.""" + from automem.consolidation.entity_dedup import merge_entities + + graph = EntityFakeGraph() + graph.entities["entity:people:alice-smith"] = { + "id": "entity:people:alice-smith", + "slug": "alice-smith", + "category": "people", + "aliases": [], + "merged_into": None, + } + graph.entities["entity:people:alice"] = { + "id": "entity:people:alice", + "slug": "alice", + "category": "people", + "aliases": [], + "merged_into": None, + } + graph.memories["m1"] = {"id": "m1", "content": "test", "tags": []} + graph.entity_edges.append({"entity_id": "entity:people:alice", "memory_id": "m1"}) + + result = merge_entities(graph, "entity:people:alice-smith", "entity:people:alice") + + assert result.canonical_id == "entity:people:alice-smith" + assert result.alias_slug == "alice" + assert result.edges_moved == 1 + assert graph.entities["entity:people:alice"]["merged_into"] == "entity:people:alice-smith" + assert "alice" in graph.entities["entity:people:alice-smith"]["aliases"] + + +# --------------------------------------------------------------------------- +# Test: Identity synthesis +# --------------------------------------------------------------------------- + + +class TestIdentitySynthesis: + """Test identity synthesis with mocked LLM.""" + + def test_synthesize_identity(self) -> None: + """Identity synthesis calls LLM and stores result.""" + from automem.consolidation.identity_synthesis import synthesize_identity + + graph = EntityFakeGraph() + graph.entities["entity:people:alice-smith"] = { + "id": "entity:people:alice-smith", + "slug": "alice-smith", + "category": "people", + "name": "Alice Smith", + "aliases": [], + "identity": None, + "identity_version": 0, + "identity_updated_at": None, + "identity_source_count": 3, + "merged_into": None, + } + graph.memories["m1"] = { + "id": "m1", + "content": "Alice works at Acme Corp as a manager.", + "importance": 0.9, + "timestamp": "2026-01-15T12:00:00Z", + "type": "Context", + "tags": ["entity:people:alice-smith"], + } + graph.entity_edges.append({"entity_id": "entity:people:alice-smith", "memory_id": "m1"}) + + # Mock OpenAI client + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content="Alice Smith is a manager at Acme Corp.") + ) + ] + ) + + result = synthesize_identity( + graph, "entity:people:alice-smith", mock_client, model="test-model" + ) + + assert result == "Alice Smith is a manager at Acme Corp." + assert graph.entities["entity:people:alice-smith"]["identity"] == result + assert graph.entities["entity:people:alice-smith"]["identity_version"] == 1 + mock_client.chat.completions.create.assert_called_once() + + def test_synthesize_no_memories(self) -> None: + """Returns None when entity has no linked memories.""" + from automem.consolidation.identity_synthesis import synthesize_identity + + graph = EntityFakeGraph() + graph.entities["entity:people:unknown"] = { + "id": "entity:people:unknown", + "slug": "unknown", + "category": "people", + "name": "Unknown", + "aliases": [], + "identity": None, + "identity_version": 0, + "identity_updated_at": None, + "identity_source_count": 0, + "merged_into": None, + } + + mock_client = MagicMock() + result = synthesize_identity(graph, "entity:people:unknown", mock_client) + + assert result is None + mock_client.chat.completions.create.assert_not_called() + + def test_run_identity_consolidation(self) -> None: + """Full identity consolidation run with dedup + synthesis.""" + from automem.consolidation.identity_synthesis import run_identity_consolidation + + graph = EntityFakeGraph() + graph.entities["entity:people:alice-smith"] = { + "id": "entity:people:alice-smith", + "slug": "alice-smith", + "category": "people", + "name": "Alice Smith", + "aliases": [], + "identity": None, + "identity_version": 0, + "identity_updated_at": None, + "identity_source_count": 3, + "merged_into": None, + } + graph.memories["m1"] = { + "id": "m1", + "content": "Alice works at Acme Corp.", + "importance": 0.9, + "timestamp": "2026-01-15T12:00:00Z", + "type": "Context", + "tags": ["entity:people:alice-smith"], + } + graph.entity_edges.append({"entity_id": "entity:people:alice-smith", "memory_id": "m1"}) + + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = SimpleNamespace( + choices=[ + SimpleNamespace(message=SimpleNamespace(content="Alice is a manager at Acme Corp.")) + ] + ) + + result = run_identity_consolidation(graph, mock_client, min_references=1) + + assert result["identities_synthesized"] == 1 + assert result["entities_examined"] >= 1 + + def test_dry_run_no_llm_calls(self) -> None: + """Dry run counts but doesn't call LLM.""" + from automem.consolidation.identity_synthesis import run_identity_consolidation + + graph = EntityFakeGraph() + graph.entities["entity:people:alice-smith"] = { + "id": "entity:people:alice-smith", + "slug": "alice-smith", + "category": "people", + "name": "Alice Smith", + "aliases": [], + "identity": None, + "identity_version": 0, + "identity_updated_at": None, + "identity_source_count": 3, + "merged_into": None, + } + # Add an edge so ref_count > 0 + graph.memories["m1"] = { + "id": "m1", + "content": "test", + "tags": [], + "importance": 0.5, + } + graph.entity_edges.append({"entity_id": "entity:people:alice-smith", "memory_id": "m1"}) + + mock_client = MagicMock() + result = run_identity_consolidation(graph, mock_client, dry_run=True) + + assert result["identities_synthesized"] == 1 # counted but not actually called + mock_client.chat.completions.create.assert_not_called() + + def test_skip_unchanged_entities(self) -> None: + """Entities with existing identity and matching source count are skipped.""" + from automem.consolidation.identity_synthesis import run_identity_consolidation + + graph = EntityFakeGraph() + graph.entities["entity:people:alice-smith"] = { + "id": "entity:people:alice-smith", + "slug": "alice-smith", + "category": "people", + "name": "Alice Smith", + "aliases": [], + "identity": "Alice is a manager at Acme Corp.", + "identity_version": 1, + "identity_updated_at": "2026-01-15T12:00:00Z", + "identity_source_count": 1, # matches actual edge count + "merged_into": None, + } + graph.memories["m1"] = { + "id": "m1", + "content": "Alice works at Acme Corp.", + "importance": 0.9, + "timestamp": "2026-01-15T12:00:00Z", + "type": "Context", + "tags": ["entity:people:alice-smith"], + } + graph.entity_edges.append({"entity_id": "entity:people:alice-smith", "memory_id": "m1"}) + + mock_client = MagicMock() + result = run_identity_consolidation(graph, mock_client, min_references=1) + + # Should skip — identity exists and source count matches + assert result["identities_synthesized"] == 0 + mock_client.chat.completions.create.assert_not_called() + + +# --------------------------------------------------------------------------- +# Test: Recall with entity injection +# --------------------------------------------------------------------------- + + +class TestRecallEntityInjection: + """Test entity identity injection in recall response.""" + + def test_extract_query_entities(self) -> None: + """Query entity extraction picks up capitalized names.""" + from automem.api.recall import _extract_query_entities + + entities = _extract_query_entities("Tell me about Alice and her work at Acme Corp") + # "Alice" is first word (skipped), but "Acme" or "Corp" should be extracted + assert "Acme" in entities or "Corp" in entities + + def test_extract_entities_from_results(self) -> None: + """Entity tags in results are extracted.""" + from automem.api.recall import _extract_entities_from_results + + results = [ + { + "memory": { + "tags": ["entity:people:alice-smith", "work"], + "metadata": {}, + } + } + ] + entities = _extract_entities_from_results(results) + assert "alice smith" in entities or "alice-smith" in {e.replace(" ", "-") for e in entities}