@dreamtree-org/graphify 1.3.0 → 1.4.0

package/dist/cli/index.js

@@ -13,7 +13,7 @@ import {
   runPipeline,
   saveResult,
   validateRules
-} from "../chunk-UBXYMMXJ.js";
+} from "../chunk-QEB7A5KB.js";
 import {
   affectedBy,
   buildContextPack,
@@ -25,7 +25,7 @@ import {
   shortestPath,
   testsForChangedFiles,
   testsForNode
-} from "../chunk-K322XB4U.js";
+} from "../chunk-7LTO76UD.js";
 import {
   validateDsn,
   validateUrl

package/dist/index.cjs

@@ -44,6 +44,7 @@ __export(index_exports, {
   collectFiles: () => collectFiles,
   deserializeGraph: () => deserializeGraph,
   diffGraphs: () => diffGraphs,
+  expandRetrieval: () => expandRetrieval,
   explainNode: () => explainNode,
   exportGraph: () => exportGraph,
   extract: () => extract,
@@ -3200,6 +3201,141 @@ function explainNode(graph, query) {
   };
 }
 
+// src/retrieval.ts
+var DEFAULT_TOKEN_BUDGET2 = 2e3;
+var DEFAULT_MAX_HOPS = 2;
+function resolveSeeds(graph, hits) {
+  const seeds = /* @__PURE__ */ new Map();
+  for (const hit of hits) {
+    let id = null;
+    if (hit.nodeId !== void 0 && graph.hasNode(hit.nodeId)) {
+      id = hit.nodeId;
+    } else if (hit.text !== void 0) {
+      id = resolveNode(graph, hit.text)?.id ?? null;
+    }
+    const key = id ?? hit.nodeId ?? hit.text ?? "";
+    if (key === "") continue;
+    const existing = seeds.get(key);
+    if (existing) {
+      if (hit.score !== void 0 && (existing.score === void 0 || hit.score > existing.score)) {
+        existing.score = hit.score;
+      }
+      continue;
+    }
+    seeds.set(key, { id: key, score: hit.score, inGraph: id !== null, text: hit.text });
+  }
+  return [...seeds.values()];
+}
+function expandRetrieval(graph, hits, options = {}) {
+  const tokenBudget = options.tokenBudget ?? DEFAULT_TOKEN_BUDGET2;
+  const maxHops = options.maxHops ?? DEFAULT_MAX_HOPS;
+  const relations = options.relations ? new Set(options.relations) : null;
+  const seeds = resolveSeeds(graph, hits);
+  const graphSeeds = seeds.filter((s) => s.inGraph);
+  const seedScore = new Map(graphSeeds.map((s) => [s.id, s.score]));
+  const hops = /* @__PURE__ */ new Map();
+  const via = /* @__PURE__ */ new Map();
+  const frontier = [];
+  let spentTokens = 0;
+  for (const seed of graphSeeds) {
+    hops.set(seed.id, 0);
+    frontier.push(seed.id);
+    spentTokens += nodeTokenCost(graph, seed.id);
+  }
+  let truncated = false;
+  while (frontier.length > 0) {
+    const current = frontier.shift();
+    const currentHops = hops.get(current);
+    if (currentHops >= maxHops) continue;
+    const candidates = [];
+    graph.forEachEdge(current, (_key, attrs, source, target) => {
+      const relation = String(attrs.relation);
+      if (relations && !relations.has(relation)) return;
+      const neighbor = source === current ? target : source;
+      candidates.push({
+        neighbor,
+        edge: { source, target, relation, confidence: String(attrs.confidence) }
+      });
+    });
+    candidates.sort((a, b) => a.neighbor.localeCompare(b.neighbor));
+    for (const { neighbor, edge } of candidates) {
+      if (hops.has(neighbor)) continue;
+      const cost = nodeTokenCost(graph, neighbor);
+      if (spentTokens + cost > tokenBudget) {
+        truncated = true;
+        continue;
+      }
+      spentTokens += cost;
+      hops.set(neighbor, currentHops + 1);
+      via.set(neighbor, edge);
+      frontier.push(neighbor);
+    }
+  }
+  const degree = (id) => graph.degree(id);
+  const nodes = [...hops.entries()].map(([id, hopCount]) => {
+    const attrs = graph.getNodeAttributes(id);
+    return {
+      id,
+      label: attrs.label ?? id,
+      ...attrs.kind !== void 0 ? { kind: String(attrs.kind) } : {},
+      sourceFile: attrs.sourceFile ?? "",
+      sourceLocation: attrs.sourceLocation ?? "",
+      hops: hopCount,
+      ...via.has(id) ? { via: via.get(id) } : {},
+      ...hopCount === 0 && seedScore.get(id) !== void 0 ? { seedScore: seedScore.get(id) } : {}
+    };
+  });
+  for (const seed of seeds) {
+    if (seed.inGraph) continue;
+    nodes.push({
+      id: seed.id,
+      label: seed.text ?? seed.id,
+      sourceFile: "",
+      sourceLocation: "",
+      hops: 0,
+      ...seed.score !== void 0 ? { seedScore: seed.score } : {}
+    });
+  }
+  nodes.sort((a, b) => {
+    if (a.hops !== b.hops) return a.hops - b.hops;
+    if (a.hops === 0) {
+      return (b.seedScore ?? -Infinity) - (a.seedScore ?? -Infinity) || a.id.localeCompare(b.id);
+    }
+    const degreeDiff = (graph.hasNode(b.id) ? degree(b.id) : 0) - (graph.hasNode(a.id) ? degree(a.id) : 0);
+    return degreeDiff || a.id.localeCompare(b.id);
+  });
+  const included = /* @__PURE__ */ new Set([...hops.keys()]);
+  const edges = [];
+  graph.forEachEdge((_key, attrs, source, target) => {
+    if (!included.has(source) || !included.has(target)) return;
+    const relation = String(attrs.relation);
+    if (relations && !relations.has(relation)) return;
+    edges.push({ source, target, relation, confidence: String(attrs.confidence) });
+  });
+  edges.sort(
+    (a, b) => a.source.localeCompare(b.source) || a.target.localeCompare(b.target) || a.relation.localeCompare(b.relation)
+  );
+  const communities = /* @__PURE__ */ new Map();
+  for (const id of included) {
+    const attrs = graph.getNodeAttributes(id);
+    const community = attrs.community;
+    if (community === void 0) continue;
+    const entry = communities.get(community) ?? {
+      id: community,
+      label: attrs.communityLabel ?? String(community),
+      seedCount: 0
+    };
+    if (hops.get(id) === 0) entry.seedCount += 1;
+    communities.set(community, entry);
+  }
+  return {
+    nodes,
+    edges,
+    communities: [...communities.values()].sort((a, b) => b.seedCount - a.seedCount || a.id - b.id),
+    truncated
+  };
+}
+
 // src/impact.ts
 var DEPENDENCY_RELATIONS = /* @__PURE__ */ new Set([
   "calls",
@@ -4005,6 +4141,7 @@ function checkRules(graph, config) {
   collectFiles,
   deserializeGraph,
   diffGraphs,
+  expandRetrieval,
   explainNode,
   exportGraph,
   extract,