squish-memory 1.1.5 → 1.2.0

package/core/graph/entity-deduplicator.ts

@@ -0,0 +1,232 @@
+/**
+ * Entity Deduplicator
+ * 
+ * Deduplicates entities using embedding similarity and string matching.
+ * Critical for knowledge graph - prevents duplicate nodes for aliases like "AWS" / "Amazon Web Services".
+ */
+
+import { eq, and, or, sql } from 'drizzle-orm';
+import { getDb } from '../../db/index.js';
+import { getSchema } from '../../db/schema.js';
+import { getEmbedding } from '../embeddings.js';
+import { cosineSimilarity } from '../utils/vector-operations.js';
+import { parseEmbedding } from '../lib/parse-embedding.js';
+import { logger } from '../logger.js';
+
+export interface DeduplicationResult {
+  merged: number;
+  aliases: Array<{ from: string; to: string; similarity: number }>;
+  totalEntities: number;
+  uniqueEntities: number;
+}
+
+// Common entity aliases for normalization
+const KNOWN_ALIASES: Record<string, string[]> = {
+  'amazon web services': ['aws'],
+  'aws': ['amazon web services'],
+  'postgresql': ['postgres', 'pg'],
+  'kubernetes': ['k8s'],
+  'k8s': ['kubernetes'],
+  'javascript': ['js'],
+  'typescript': ['ts'],
+  'node.js': ['nodejs', 'node'],
+  'react.js': ['react'],
+  'vue.js': ['vue'],
+  'machine learning': ['ml'],
+  'artificial intelligence': ['ai'],
+  'large language model': ['llm'],
+  'natural language processing': ['nlp'],
+};
+
+/**
+ * Find and merge duplicate entities in a project.
+ */
+export async function deduplicateProjectEntities(
+  projectId: string,
+  options?: {
+    similarityThreshold?: number;
+    dryRun?: boolean;
+  }
+): Promise<DeduplicationResult> {
+  const { similarityThreshold = 0.85, dryRun = false } = options || {};
+  const db = await getDb();
+  const schema = await getSchema();
+
+  // Get all entities for the project
+  const entities = await (db as any)
+    .select()
+    .from(schema.entities)
+    .where(eq(schema.entities.projectId, projectId));
+
+  if (entities.length < 2) {
+    return {
+      merged: 0,
+      aliases: [],
+      totalEntities: entities.length,
+      uniqueEntities: entities.length,
+    };
+  }
+
+  const aliases: Array<{ from: string; to: string; similarity: number }> = [];
+  const merged = new Set<string>();
+  let mergeCount = 0;
+
+  // Compare each pair of entities
+  for (let i = 0; i < entities.length; i++) {
+    if (merged.has(entities[i].id)) continue;
+
+    for (let j = i + 1; j < entities.length; j++) {
+      if (merged.has(entities[j].id)) continue;
+
+      const similarity = await computeEntitySimilarity(entities[i], entities[j]);
+
+      if (similarity >= similarityThreshold) {
+        // Merge j into i (keep the one with more mentions)
+        const target = (entities[i].mentionCount || 0) >= (entities[j].mentionCount || 0)
+          ? entities[i]
+          : entities[j];
+        const source = target.id === entities[i].id ? entities[j] : entities[i];
+
+        aliases.push({
+          from: source.name,
+          to: target.name,
+          similarity,
+        });
+
+        if (!dryRun) {
+          await mergeEntities(source, target, db, schema);
+        }
+
+        merged.add(source.id);
+        mergeCount++;
+      }
+    }
+  }
+
+  logger.info('Entity deduplication completed', {
+    projectId,
+    totalEntities: entities.length,
+    merged: mergeCount,
+    dryRun,
+  });
+
+  return {
+    merged: mergeCount,
+    aliases,
+    totalEntities: entities.length,
+    uniqueEntities: entities.length - mergeCount,
+  };
+}
+
+/**
+ * Compute similarity between two entities.
+ * Uses a combination of string matching, known aliases, and embedding similarity.
+ */
+async function computeEntitySimilarity(entity1: any, entity2: any): Promise<number> {
+  // Must be same type to be considered duplicates
+  if (entity1.type !== entity2.type) {
+    return 0;
+  }
+
+  const name1 = entity1.name.toLowerCase().trim();
+  const name2 = entity2.name.toLowerCase().trim();
+
+  // Exact match
+  if (name1 === name2) return 1.0;
+
+  // Check known aliases
+  const name1Aliases = KNOWN_ALIASES[name1] || [];
+  const name2Aliases = KNOWN_ALIASES[name2] || [];
+
+  if (name1Aliases.includes(name2) || name2Aliases.includes(name1)) {
+    return 0.99; // Very high confidence for known aliases
+  }
+
+  // String similarity (Jaccard on words)
+  const words1 = new Set(name1.split(/[\s._-]+/).filter((w: string) => w.length > 1));
+  const words2 = new Set(name2.split(/[\s._-]+/).filter((w: string) => w.length > 1));
+
+  const words1Arr = Array.from(words1);
+  const words2Arr = Array.from(words2);
+  const intersection = words1Arr.filter(w => words2.has(w));
+  const union = new Set([...words1Arr, ...words2Arr]);
+
+  if (words1.size > 0 && words2.size > 0) {
+    const jaccard = intersection.length / union.size;
+
+    if (jaccard >= 0.8) return jaccard;
+  }
+
+  // One contains the other (e.g., "React" in "React.js")
+  if (name1.includes(name2) || name2.includes(name1)) {
+    const shorter = name1.length < name2.length ? name1 : name2;
+    const longer = name1.length < name2.length ? name2 : name1;
+    // Only count as similar if the shorter name is a significant portion
+    if (shorter.length / longer.length >= 0.5) {
+      return 0.85;
+    }
+  }
+
+  // Embedding similarity (if available)
+  const embedding1 = parseEmbedding(entity1.embedding || entity1.embedding_json);
+  const embedding2 = parseEmbedding(entity2.embedding || entity2.embedding_json);
+
+  if (embedding1 && embedding2) {
+    const embSimilarity = cosineSimilarity(embedding1, embedding2);
+    if (embSimilarity >= similarityThreshold) {
+      return embSimilarity;
+    }
+  }
+
+  return 0;
+}
+
+/**
+ * Merge source entity into target entity.
+ * Updates all relations pointing to source to point to target instead.
+ */
+async function mergeEntities(
+  source: any,
+  target: any,
+  db: any,
+  schema: any
+): Promise<void> {
+  // Update all relations pointing FROM source to point FROM target
+  await db
+    .update(schema.entityRelations)
+    .set({ fromEntityId: target.id })
+    .where(eq(schema.entityRelations.fromEntityId, source.id));
+
+  // Update all relations pointing TO source to point TO target
+  await db
+    .update(schema.entityRelations)
+    .set({ toEntityId: target.id })
+    .where(eq(schema.entityRelations.toEntityId, source.id));
+
+  // Add source name as alias to target
+  const existingAliases = Array.isArray(target.aliases) ? target.aliases as string[] : [];
+  const updatedAliases = [...new Set([...existingAliases, source.name])];
+
+  await db
+    .update(schema.entities)
+    .set({
+      aliases: updatedAliases as any,
+      mentionCount: (target.mentionCount || 0) + (source.mentionCount || 0),
+      updatedAt: new Date(),
+    })
+    .where(eq(schema.entities.id, target.id));
+
+  // Delete source entity
+  await db
+    .delete(schema.entities)
+    .where(eq(schema.entities.id, source.id));
+
+  logger.debug('Merged entity', {
+    source: source.name,
+    target: target.name,
+    aliases: updatedAliases,
+  });
+}
+
+// Module-level threshold for embedding similarity
+const similarityThreshold = 0.85;

package/core/graph/graph-builder.ts

@@ -0,0 +1,257 @@
+/**
+ * Graph Builder
+ * 
+ * Orchestrates the full pipeline: extract entities and relations from
+ * memories and store them in the knowledge graph.
+ */
+
+import { eq } from 'drizzle-orm';
+import { getDb } from '../../db/index.js';
+import { getSchema } from '../../db/schema.js';
+import { logger } from '../logger.js';
+import { config } from '../../config.js';
+import { extractAndStoreRelations } from './relationship-extractor.js';
+import { deduplicateProjectEntities } from './entity-deduplicator.js';
+
+export interface GraphBuildStats {
+  memoriesProcessed: number;
+  entitiesCreated: number;
+  relationsCreated: number;
+  entitiesDeduplicated: number;
+  errors: number;
+  durationMs: number;
+}
+
+export interface GraphAddStats {
+  entitiesCreated: number;
+  relationsCreated: number;
+  source: 'llm' | 'regex' | 'fallback' | 'none';
+}
+
+/**
+ * Build or rebuild the entity graph for a project.
+ * Processes all memories in the project, extracting entities and relations.
+ */
+export async function buildGraphForProject(
+  projectPath: string,
+  options?: {
+    clearExisting?: boolean;
+    batchSize?: number;
+    preferLLM?: boolean;
+    deduplicate?: boolean;
+  }
+): Promise<GraphBuildStats> {
+  const startTime = Date.now();
+  const {
+    clearExisting = false,
+    batchSize = 10,
+    preferLLM = true,
+    deduplicate = true,
+  } = options || {};
+
+  const db = await getDb();
+  const schema = await getSchema();
+
+  // Get project
+  const projectRows = await (db as any)
+    .select()
+    .from(schema.projects)
+    .where(eq(schema.projects.path, projectPath))
+    .limit(1);
+
+  if (projectRows.length === 0) {
+    logger.warn('Project not found for graph build', { projectPath });
+    return {
+      memoriesProcessed: 0,
+      entitiesCreated: 0,
+      relationsCreated: 0,
+      entitiesDeduplicated: 0,
+      errors: 0,
+      durationMs: Date.now() - startTime,
+    };
+  }
+
+  const projectId = projectRows[0].id;
+
+  // Clear existing graph if requested
+  if (clearExisting) {
+    const { clearProjectGraph } = await import('./relationship-extractor.js');
+    await clearProjectGraph(projectId);
+    logger.info('Cleared existing graph for rebuild', { projectId });
+  }
+
+  // Get all active memories for the project
+  const memories = await (db as any)
+    .select()
+    .from(schema.memories)
+    .where(eq(schema.memories.projectId, projectId));
+
+  logger.info('Starting graph build for project', {
+    projectId,
+    memoryCount: memories.length,
+    clearExisting,
+  });
+
+  let entitiesCreated = 0;
+  let relationsCreated = 0;
+  let errors = 0;
+
+  // Process memories in batches
+  for (let i = 0; i < memories.length; i += batchSize) {
+    const batch = memories.slice(i, i + batchSize);
+
+    for (const memory of batch) {
+      try {
+        const result = await extractAndStoreRelations(
+          memory.id,
+          memory.content,
+          projectId,
+          { preferLLM }
+        );
+
+        entitiesCreated += result.entities;
+        relationsCreated += result.relations;
+      } catch (error) {
+        logger.error('Error processing memory for graph', {
+          memoryId: memory.id,
+          error: error as Error,
+        });
+        errors++;
+      }
+    }
+
+    logger.debug('Graph build batch progress', {
+      processed: Math.min(i + batchSize, memories.length),
+      total: memories.length,
+    });
+  }
+
+  // Deduplicate entities
+  let entitiesDeduplicated = 0;
+  if (deduplicate) {
+    try {
+      const dedupResult = await deduplicateProjectEntities(projectId);
+      entitiesDeduplicated = dedupResult.merged;
+    } catch (error) {
+      logger.error('Error during entity deduplication', { error: error as Error });
+    }
+  }
+
+  const stats: GraphBuildStats = {
+    memoriesProcessed: memories.length,
+    entitiesCreated,
+    relationsCreated,
+    entitiesDeduplicated,
+    errors,
+    durationMs: Date.now() - startTime,
+  };
+
+  logger.info('Graph build completed', stats);
+
+  return stats;
+}
+
+/**
+ * Add a single memory to the knowledge graph.
+ * Used for incremental updates when new memories are stored.
+ */
+export async function addMemoryToGraph(
+  memoryId: string,
+  options?: {
+    preferLLM?: boolean;
+  }
+): Promise<GraphAddStats> {
+  const db = await getDb();
+  const schema = await getSchema();
+
+  // Get the memory
+  const memories = await (db as any)
+    .select()
+    .from(schema.memories)
+    .where(eq(schema.memories.id, memoryId))
+    .limit(1);
+
+  if (memories.length === 0) {
+    return { entitiesCreated: 0, relationsCreated: 0, source: 'none' };
+  }
+
+  const memory = memories[0];
+
+  if (!memory.projectId) {
+    return { entitiesCreated: 0, relationsCreated: 0, source: 'none' as const };
+  }
+
+  // Use global config if not explicitly overridden
+  const preferLLM = options?.preferLLM ?? config.llmEnabled;
+
+  const result = await extractAndStoreRelations(
+    memoryId,
+    memory.content,
+    memory.projectId,
+    { preferLLM }
+  );
+
+  return {
+    entitiesCreated: result.entities,
+    relationsCreated: result.relations,
+    source: result.source,
+  };
+}
+
+/**
+ * Get graph statistics for a project.
+ */
+export async function getGraphStats(projectPath: string): Promise<{
+  entityCount: number;
+  relationCount: number;
+  relationTypes: Record<string, number>;
+  avgConnections: number;
+}> {
+  const db = await getDb();
+  const schema = await getSchema();
+
+  const projectRows = await (db as any)
+    .select({ id: schema.projects.id })
+    .from(schema.projects)
+    .where(eq(schema.projects.path, projectPath))
+    .limit(1);
+
+  if (projectRows.length === 0) {
+    return { entityCount: 0, relationCount: 0, relationTypes: {}, avgConnections: 0 };
+  }
+
+  const projectId = projectRows[0].id;
+
+  const entities = await (db as any)
+    .select()
+    .from(schema.entities)
+    .where(eq(schema.entities.projectId, projectId));
+
+  const relations = await (db as any)
+    .select()
+    .from(schema.entityRelations);
+
+  // Filter relations to only those involving project entities
+  const entityIds = new Set(entities.map((e: any) => e.id));
+  const projectRelations = relations.filter(
+    (r: any) => entityIds.has(r.fromEntityId) || entityIds.has(r.toEntityId)
+  );
+
+  // Count by type
+  const relationTypes: Record<string, number> = {};
+  for (const r of projectRelations) {
+    const type = r.type || 'unknown';
+    relationTypes[type] = (relationTypes[type] || 0) + 1;
+  }
+
+  const avgConnections = entities.length > 0
+    ? (projectRelations.length * 2) / entities.length
+    : 0;
+
+  return {
+    entityCount: entities.length,
+    relationCount: projectRelations.length,
+    relationTypes,
+    avgConnections: Math.round(avgConnections * 100) / 100,
+  };
+}