@devilsdev/rag-pipeline-utils 2.3.0 → 2.4.0

package/src/graph/graph-index.js

@@ -0,0 +1,113 @@
+'use strict';
+
+/**
+ * Fast entity lookup index for the knowledge graph.
+ *
+ * Maintains secondary indexes by type, name, and arbitrary properties
+ * for efficient querying without full graph traversal.
+ */
+class GraphIndex {
+  constructor() {
+    /** @type {Map<string, Set<string>>} type -> Set of entity IDs */
+    this.byType = new Map();
+    /** @type {Map<string, string>} normalized name -> entity ID */
+    this.byName = new Map();
+    /** @type {Map<string, Map<*, Set<string>>>} propertyKey -> Map of value -> Set of entity IDs */
+    this.byProperty = new Map();
+  }
+
+  /**
+   * Index an entity by type, name, and each of its properties.
+   * @param {{id: string, type: string, properties: Object}} entity
+   */
+  indexEntity(entity) {
+    const { id, type, properties } = entity;
+
+    // Index by type
+    if (!this.byType.has(type)) {
+      this.byType.set(type, new Set());
+    }
+    this.byType.get(type).add(id);
+
+    // Index by normalized name
+    const name = (properties && properties.name) || entity.name || id;
+    const normalized = String(name).toLowerCase();
+    this.byName.set(normalized, id);
+
+    // Index by each property
+    if (properties) {
+      for (const [key, value] of Object.entries(properties)) {
+        if (!this.byProperty.has(key)) {
+          this.byProperty.set(key, new Map());
+        }
+        const valueMap = this.byProperty.get(key);
+        if (!valueMap.has(value)) {
+          valueMap.set(value, new Set());
+        }
+        valueMap.get(value).add(id);
+      }
+    }
+  }
+
+  /**
+   * Fuzzy search for entities whose name contains the query (case-insensitive).
+   *
+   * Results are sorted by relevance: exact matches first, then partial matches.
+   *
+   * @param {string} query - Search query
+   * @returns {string[]} Array of matching entity IDs
+   */
+  searchByName(query) {
+    if (!query || typeof query !== 'string') {
+      return [];
+    }
+
+    const normalized = query.toLowerCase();
+    const exact = [];
+    const partial = [];
+
+    for (const [name, entityId] of this.byName.entries()) {
+      if (name === normalized) {
+        exact.push(entityId);
+      } else if (name.includes(normalized)) {
+        partial.push(entityId);
+      }
+    }
+
+    return [...exact, ...partial];
+  }
+
+  /**
+   * Return all entity IDs of a given type.
+   * @param {string} type - Entity type
+   * @returns {string[]} Array of entity IDs
+   */
+  searchByType(type) {
+    const ids = this.byType.get(type);
+    return ids ? [...ids] : [];
+  }
+
+  /**
+   * Return entity IDs matching a property key-value pair.
+   * @param {string} key - Property key
+   * @param {*} value - Property value
+   * @returns {string[]} Array of entity IDs
+   */
+  searchByProperty(key, value) {
+    const valueMap = this.byProperty.get(key);
+    if (!valueMap) return [];
+    const ids = valueMap.get(value);
+    return ids ? [...ids] : [];
+  }
+
+  /**
+   * Clear all indexes.
+   */
+  clear() {
+    this.byType.clear();
+    this.byName.clear();
+    this.byProperty.clear();
+  }
+}
+
+module.exports = { GraphIndex };

package/src/graph/graph-retriever.js

@@ -0,0 +1,244 @@
+'use strict';
+
+const { EventEmitter } = require('events');
+const { KnowledgeGraph } = require('./knowledge-graph');
+const { EntityExtractor } = require('./entity-extractor');
+const { GraphIndex } = require('./graph-index');
+const { tokenize } = require('../evaluate/scoring');
+
+/**
+ * Default configuration for the GraphRetriever.
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  traversalDepth: 2,
+  maxResults: 10,
+  entityWeight: 0.6,
+  textWeight: 0.4,
+};
+
+/**
+ * Retriever that uses a knowledge graph for document retrieval.
+ *
+ * Combines entity-based graph traversal with token-overlap text scoring
+ * to rank and retrieve relevant documents from a corpus.
+ */
+class GraphRetriever extends EventEmitter {
+  /**
+   * @param {Object} [options={}] - Retriever options
+   * @param {number} [options.traversalDepth=2] - Graph traversal depth
+   * @param {number} [options.maxResults=10] - Maximum results to return
+   * @param {number} [options.entityWeight=0.6] - Weight for entity-graph score
+   * @param {number} [options.textWeight=0.4] - Weight for text-overlap score
+   * @param {Object} [options.extractorOptions] - Options for the EntityExtractor
+   */
+  constructor(options = {}) {
+    super();
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    this.graph = new KnowledgeGraph();
+    this.extractor = new EntityExtractor(options.extractorOptions);
+    this.index = new GraphIndex();
+    /** @type {Map<string, {id: string, content: string, metadata: Object}>} */
+    this.documents = new Map();
+  }
+
+  /**
+   * Store documents, extract entities and relationships, and build the graph.
+   *
+   * @param {Array<{id: string, content: string, metadata?: Object}>} documents
+   * @returns {Promise<void>}
+   */
+  async store(documents) {
+    if (!Array.isArray(documents)) return;
+
+    for (const doc of documents) {
+      const id = doc.id || `doc_${this.documents.size + 1}`;
+      this.documents.set(id, {
+        id,
+        content: doc.content || '',
+        metadata: doc.metadata || {},
+      });
+    }
+
+    const { entities, relationships } = this.extractor.extractFromDocuments(
+      documents.map((doc) => ({
+        id: doc.id || `doc_${this.documents.size}`,
+        content: doc.content || '',
+        metadata: doc.metadata || {},
+      })),
+    );
+
+    for (const entity of entities) {
+      const entityId = entity.name.toLowerCase().trim().replace(/\s+/g, '_');
+      this.graph.addEntity(entityId, entity.type, {
+        name: entity.name,
+        frequency: entity.frequency,
+        sources: entity.sources || [],
+      });
+      this.index.indexEntity({
+        id: entityId,
+        type: entity.type,
+        name: entity.name,
+        properties: {
+          name: entity.name,
+          frequency: entity.frequency,
+        },
+      });
+    }
+
+    for (const rel of relationships) {
+      const fromId = rel.from.toLowerCase().trim().replace(/\s+/g, '_');
+      const toId = rel.to.toLowerCase().trim().replace(/\s+/g, '_');
+
+      if (this.graph.getEntity(fromId) && this.graph.getEntity(toId)) {
+        this.graph.addRelationship(fromId, toId, rel.type, {
+          context: rel.context,
+          source: rel.source,
+        });
+      }
+    }
+
+    this.emit('stored', {
+      documentCount: documents.length,
+      stats: this.graph.getStats(),
+    });
+  }
+
+  /**
+   * Retrieve documents relevant to a query using graph + text scoring.
+   *
+   * @param {string|{query: string, topK?: number}} queryOrObj - Query string or options object
+   * @param {number} [k] - Number of results (overrides config.maxResults)
+   * @returns {Promise<Array<{id: string, score: number, content: string, metadata: Object}>>}
+   */
+  async retrieve(queryOrObj, k) {
+    const query =
+      typeof queryOrObj === 'string' ? queryOrObj : queryOrObj.query;
+    const topK =
+      k ||
+      (typeof queryOrObj === 'object' && queryOrObj.topK) ||
+      this.config.maxResults;
+
+    // Extract entities from query
+    const queryEntities = this.extractor.extract(query);
+    const queryTokens = tokenize(query);
+
+    // Score documents
+    /** @type {Map<string, {graphScore: number, textScore: number, entities: string[], relationships: Array}>} */
+    const docScores = new Map();
+
+    // Initialize all documents with zero scores
+    for (const [docId] of this.documents) {
+      docScores.set(docId, {
+        graphScore: 0,
+        textScore: 0,
+        entities: [],
+        relationships: [],
+      });
+    }
+
+    // Graph-based scoring: find matching entities and traverse
+    for (const qEntity of queryEntities) {
+      const normalized = qEntity.name.toLowerCase().trim();
+      const matchingIds = this.index.searchByName(normalized);
+
+      for (const entityId of matchingIds) {
+        const traversal = this.graph.traverse(
+          entityId,
+          this.config.traversalDepth,
+        );
+
+        for (const entity of traversal.entities) {
+          const sources =
+            (entity.properties && entity.properties.sources) || [];
+          for (const sourceDocId of sources) {
+            if (docScores.has(sourceDocId)) {
+              const scores = docScores.get(sourceDocId);
+              scores.graphScore += 1;
+              if (!scores.entities.includes(entity.id)) {
+                scores.entities.push(entity.id);
+              }
+            }
+          }
+        }
+
+        for (const rel of traversal.relationships) {
+          for (const [docId, scores] of docScores) {
+            const docEntities = scores.entities;
+            if (
+              docEntities.includes(rel.from) ||
+              docEntities.includes(rel.to)
+            ) {
+              scores.relationships.push(rel);
+            }
+          }
+        }
+      }
+    }
+
+    // Text-based scoring: token overlap
+    if (queryTokens.length > 0) {
+      for (const [docId, doc] of this.documents) {
+        const docTokens = tokenize(doc.content);
+        if (docTokens.length === 0) continue;
+
+        const querySet = new Set(queryTokens);
+        let overlap = 0;
+        for (const token of docTokens) {
+          if (querySet.has(token)) {
+            overlap += 1;
+          }
+        }
+        const textScore = overlap / Math.max(queryTokens.length, 1);
+
+        if (docScores.has(docId)) {
+          docScores.get(docId).textScore = textScore;
+        }
+      }
+    }
+
+    // Combine scores and sort
+    const results = [];
+    for (const [docId, scores] of docScores) {
+      const doc = this.documents.get(docId);
+      const combinedScore =
+        this.config.entityWeight * scores.graphScore +
+        this.config.textWeight * scores.textScore;
+
+      if (combinedScore > 0) {
+        results.push({
+          id: docId,
+          score: combinedScore,
+          content: doc.content,
+          metadata: {
+            ...doc.metadata,
+            entities: scores.entities,
+            relationships: scores.relationships,
+          },
+        });
+      }
+    }
+
+    results.sort((a, b) => b.score - a.score);
+    return results.slice(0, topK);
+  }
+
+  /**
+   * Return basic graph statistics.
+   * @returns {{entityCount: number, relationshipCount: number}}
+   */
+  getGraphStats() {
+    return this.graph.getStats();
+  }
+
+  /**
+   * Clear all stored data, the graph, and indexes.
+   */
+  clear() {
+    this.graph.clear();
+    this.index.clear();
+    this.documents.clear();
+  }
+}
+
+module.exports = { GraphRetriever, DEFAULT_CONFIG };

package/src/graph/knowledge-graph.js

@@ -0,0 +1,256 @@
+'use strict';
+
+const { EventEmitter } = require('events');
+
+/**
+ * In-memory knowledge graph with entities and relationships.
+ *
+ * Stores typed entities and directed relationships between them,
+ * supporting traversal, pattern-based querying, and adjacency lookups.
+ */
+class KnowledgeGraph extends EventEmitter {
+  /**
+   * @param {Object} [options={}] - Graph configuration options
+   */
+  constructor(options = {}) {
+    super();
+    /** @type {Map<string, {id: string, type: string, properties: Object, labels: string[]}>} */
+    this.entities = new Map();
+    /** @type {Map<string, {id: string, from: string, to: string, type: string, properties: Object}>} */
+    this.relationships = new Map();
+    /** @type {Map<string, Set<string>>} */
+    this.adjacency = new Map();
+    this.options = options;
+    this._relationshipCounter = 0;
+  }
+
+  /**
+   * Add an entity to the graph.
+   * @param {string} id - Unique entity identifier
+   * @param {string} type - Entity type (e.g. 'person', 'date', 'url')
+   * @param {Object} [properties={}] - Arbitrary key-value properties
+   * @returns {{id: string, type: string, properties: Object, labels: string[]}}
+   */
+  addEntity(id, type, properties = {}) {
+    const entity = {
+      id,
+      type,
+      properties: { ...properties },
+      labels: properties.labels || [type],
+    };
+    this.entities.set(id, entity);
+
+    if (!this.adjacency.has(id)) {
+      this.adjacency.set(id, new Set());
+    }
+
+    this.emit('entityAdded', entity);
+    return entity;
+  }
+
+  /**
+   * Add a directed relationship between two entities.
+   * @param {string} fromId - Source entity ID
+   * @param {string} toId - Target entity ID
+   * @param {string} type - Relationship type (e.g. 'co-occurs', 'references')
+   * @param {Object} [properties={}] - Arbitrary key-value properties
+   * @returns {{id: string, from: string, to: string, type: string, properties: Object}}
+   * @throws {Error} If either entity does not exist
+   */
+  addRelationship(fromId, toId, type, properties = {}) {
+    if (!this.entities.has(fromId)) {
+      throw new Error(`Entity not found: ${fromId}`);
+    }
+    if (!this.entities.has(toId)) {
+      throw new Error(`Entity not found: ${toId}`);
+    }
+
+    this._relationshipCounter += 1;
+    const id = `rel_${this._relationshipCounter}`;
+
+    const relationship = {
+      id,
+      from: fromId,
+      to: toId,
+      type,
+      properties: { ...properties },
+    };
+
+    this.relationships.set(id, relationship);
+
+    if (!this.adjacency.has(fromId)) {
+      this.adjacency.set(fromId, new Set());
+    }
+    if (!this.adjacency.has(toId)) {
+      this.adjacency.set(toId, new Set());
+    }
+    this.adjacency.get(fromId).add(id);
+    this.adjacency.get(toId).add(id);
+
+    this.emit('relationshipAdded', relationship);
+    return relationship;
+  }
+
+  /**
+   * Retrieve an entity by ID.
+   * @param {string} id - Entity ID
+   * @returns {{id: string, type: string, properties: Object, labels: string[]}|null}
+   */
+  getEntity(id) {
+    return this.entities.get(id) || null;
+  }
+
+  /**
+   * Get relationships for an entity, optionally filtered.
+   * @param {string} entityId - Entity ID
+   * @param {Object} [options={}] - Filter options
+   * @param {string} [options.type] - Filter by relationship type
+   * @param {'outgoing'|'incoming'|'both'} [options.direction='both'] - Direction filter
+   * @returns {Array<{id: string, from: string, to: string, type: string, properties: Object}>}
+   */
+  getRelationships(entityId, options = {}) {
+    const { type, direction = 'both' } = options;
+    const relIds = this.adjacency.get(entityId);
+    if (!relIds) {
+      return [];
+    }
+
+    const results = [];
+    for (const relId of relIds) {
+      const rel = this.relationships.get(relId);
+      if (!rel) continue;
+
+      if (type && rel.type !== type) continue;
+
+      if (direction === 'outgoing' && rel.from !== entityId) continue;
+      if (direction === 'incoming' && rel.to !== entityId) continue;
+
+      results.push(rel);
+    }
+    return results;
+  }
+
+  /**
+   * BFS traversal from startId up to a given depth.
+   * @param {string} startId - Starting entity ID
+   * @param {number} [depth=2] - Maximum traversal depth
+   * @returns {{entities: Array, relationships: Array, paths: Array<string[]>}}
+   */
+  traverse(startId, depth = 2) {
+    const visited = new Set();
+    const resultEntities = [];
+    const resultRelationships = [];
+    const paths = [];
+
+    /** @type {Array<{id: string, depth: number, path: string[]}>} */
+    const queue = [{ id: startId, depth: 0, path: [startId] }];
+    visited.add(startId);
+
+    const startEntity = this.entities.get(startId);
+    if (startEntity) {
+      resultEntities.push(startEntity);
+    }
+
+    while (queue.length > 0) {
+      const { id, depth: currentDepth, path } = queue.shift();
+
+      if (currentDepth >= depth) {
+        paths.push(path);
+        continue;
+      }
+
+      const relIds = this.adjacency.get(id);
+      if (!relIds || relIds.size === 0) {
+        paths.push(path);
+        continue;
+      }
+
+      let expanded = false;
+      for (const relId of relIds) {
+        const rel = this.relationships.get(relId);
+        if (!rel) continue;
+
+        const neighborId = rel.from === id ? rel.to : rel.from;
+        if (visited.has(neighborId)) continue;
+
+        visited.add(neighborId);
+        expanded = true;
+
+        const neighborEntity = this.entities.get(neighborId);
+        if (neighborEntity) {
+          resultEntities.push(neighborEntity);
+        }
+        resultRelationships.push(rel);
+
+        const newPath = [...path, neighborId];
+        queue.push({ id: neighborId, depth: currentDepth + 1, path: newPath });
+      }
+
+      if (!expanded) {
+        paths.push(path);
+      }
+    }
+
+    return {
+      entities: resultEntities,
+      relationships: resultRelationships,
+      paths,
+    };
+  }
+
+  /**
+   * Query the graph by pattern matching.
+   * @param {Object} pattern - Query pattern
+   * @param {string} [pattern.entityType] - Filter entities by type
+   * @param {string} [pattern.relationshipType] - Filter relationships by type
+   * @param {Object} [pattern.properties] - Filter by property key-value pairs
+   * @returns {{entities: Array, relationships: Array}}
+   */
+  query(pattern) {
+    const { entityType, relationshipType, properties } = pattern || {};
+
+    let matchedEntities = [...this.entities.values()];
+    if (entityType) {
+      matchedEntities = matchedEntities.filter((e) => e.type === entityType);
+    }
+    if (properties) {
+      matchedEntities = matchedEntities.filter((e) =>
+        Object.entries(properties).every(
+          ([key, value]) => e.properties[key] === value,
+        ),
+      );
+    }
+
+    let matchedRelationships = [...this.relationships.values()];
+    if (relationshipType) {
+      matchedRelationships = matchedRelationships.filter(
+        (r) => r.type === relationshipType,
+      );
+    }
+
+    return { entities: matchedEntities, relationships: matchedRelationships };
+  }
+
+  /**
+   * Return basic graph statistics.
+   * @returns {{entityCount: number, relationshipCount: number}}
+   */
+  getStats() {
+    return {
+      entityCount: this.entities.size,
+      relationshipCount: this.relationships.size,
+    };
+  }
+
+  /**
+   * Clear all graph data.
+   */
+  clear() {
+    this.entities.clear();
+    this.relationships.clear();
+    this.adjacency.clear();
+    this._relationshipCounter = 0;
+  }
+}
+
+module.exports = { KnowledgeGraph };