init

ts/packages/flow/src/retrieval/graph-rag.ts

@@ -124,26 +124,156 @@ export class GraphRag {
   }
 
   private async followEdges(entities: Term[]): Promise<Triple[]> {
-    // Batch triple queries for all entities
-    const allTriples: Triple[] = [];
+    // BFS multi-hop traversal up to maxPathLength
+    const visited = new Set<string>();
+    const subgraph: Triple[] = [];
 
-    const queries = entities.map((entity) =>
-      this.clients.triples.request({ s: entity, limit: this.config.tripleLimit }),
+    // Current frontier: the set of entities to expand at this depth level
+    let currentLevel = new Set<string>(
+      entities.map((e) => termToString(e)),
     );
 
-    const results = await Promise.all(queries);
-    for (const result of results) {
-      allTriples.push(...(result as TriplesQueryResponse).triples);
+    for (let depth = 0; depth < this.config.maxPathLength; depth++) {
+      if (currentLevel.size === 0 || subgraph.length >= this.config.maxSubgraphSize) {
+        break;
+      }
+
+      // Filter out already-visited entities
+      const unvisited = [...currentLevel].filter((e) => !visited.has(e));
+      if (unvisited.length === 0) break;
+
+      // Batch triple queries for all unvisited entities at this depth
+      // Query each entity as subject to get outgoing edges
+      const queries = unvisited.map((entityStr) => {
+        const term = stringToTerm(entityStr);
+        return this.clients.triples.request({
+          s: term,
+          limit: this.config.tripleLimit,
+        });
+      });
+
+      const results = await Promise.all(queries);
+
+      const nextLevel = new Set<string>();
+
+      for (const result of results) {
+        const triples = (result as TriplesQueryResponse).triples;
+        for (const triple of triples) {
+          subgraph.push(triple);
+
+          // Collect objects as next-level entities for further expansion
+          // (only if we have more depth levels remaining)
+          if (depth < this.config.maxPathLength - 1) {
+            const objStr = termToString(triple.o);
+            if (!visited.has(objStr)) {
+              nextLevel.add(objStr);
+            }
+          }
+
+          if (subgraph.length >= this.config.maxSubgraphSize) {
+            return subgraph;
+          }
+        }
+      }
+
+      // Mark current level as visited and move to next
+      for (const e of currentLevel) {
+        visited.add(e);
+      }
+      currentLevel = nextLevel;
     }
 
-    // TODO: Multi-hop traversal up to maxPathLength
-    return allTriples.slice(0, this.config.maxSubgraphSize);
+    return subgraph.slice(0, this.config.maxSubgraphSize);
   }
 
   private async scoreEdges(query: string, triples: Triple[]): Promise<Triple[]> {
-    // TODO: LLM-based edge scoring and filtering
-    // For now, return top N edges
-    return triples.slice(0, this.config.edgeLimit);
+    if (triples.length === 0) return [];
+
+    // If the subgraph is small enough, skip LLM scoring entirely
+    if (triples.length <= this.config.edgeLimit) {
+      return triples;
+    }
+
+    // Build a numbered list of edges for the LLM to score
+    const edgeDescriptions = triples.map((t, i) => ({
+      id: String(i),
+      s: termToString(t.s),
+      p: termToString(t.p),
+      o: termToString(t.o),
+    }));
+
+    // Limit how many edges we send for scoring to avoid overflowing context
+    const toScore = edgeDescriptions.slice(0, this.config.edgeScoreLimit);
+
+    const knowledgeJson = JSON.stringify(toScore, null, 2);
+
+    // Ask the LLM to score each edge for relevance to the query
+    const promptResp = await this.clients.prompt.request({
+      name: "kg-edge-scoring",
+      variables: {
+        query,
+        knowledge: knowledgeJson,
+      },
+    });
+
+    const llmResp = await this.clients.llm.request({
+      system: (promptResp as PromptResponse).system,
+      prompt: (promptResp as PromptResponse).prompt,
+    });
+
+    const responseText = (llmResp as TextCompletionResponse).response;
+
+    // Parse scores from LLM response
+    // Expected format: JSON array of { id: string, score: number }
+    // or newline-separated JSON objects
+    const scored: Array<{ id: string; score: number }> = [];
+
+    try {
+      // Try parsing as a JSON array first
+      const parsed = JSON.parse(responseText) as Array<{ id: string; score: number }>;
+      if (Array.isArray(parsed)) {
+        for (const item of parsed) {
+          if (item && typeof item.id === "string" && typeof item.score === "number") {
+            scored.push({ id: item.id, score: item.score });
+          }
+        }
+      }
+    } catch {
+      // Fall back to parsing line-by-line JSON objects
+      for (const line of responseText.split("\n")) {
+        const trimmed = line.trim();
+        if (!trimmed) continue;
+        try {
+          const obj = JSON.parse(trimmed) as { id?: string; score?: number };
+          if (obj && typeof obj.id === "string" && typeof obj.score === "number") {
+            scored.push({ id: obj.id, score: obj.score });
+          }
+        } catch {
+          // Skip unparseable lines
+        }
+      }
+    }
+
+    // Sort by score descending and keep top N
+    scored.sort((a, b) => b.score - a.score);
+    const topN = scored.slice(0, this.config.edgeLimit);
+    const selectedIds = new Set(topN.map((e) => e.id));
+
+    // Map back to triples
+    const result: Triple[] = [];
+    for (const entry of topN) {
+      const idx = parseInt(entry.id, 10);
+      if (!isNaN(idx) && idx >= 0 && idx < triples.length) {
+        result.push(triples[idx]);
+      }
+    }
+
+    // If scoring failed entirely, fall back to returning the first edgeLimit triples
+    if (result.length === 0) {
+      return triples.slice(0, this.config.edgeLimit);
+    }
+
+    return result;
   }
 
   private async synthesize(
@@ -205,3 +335,13 @@ function termToString(term: Term): string {
       return `(${termToString(term.triple.s)} ${termToString(term.triple.p)} ${termToString(term.triple.o)})`;
   }
 }
+
+function stringToTerm(value: string): Term {
+  if (value.startsWith("http://") || value.startsWith("https://")) {
+    return { type: "IRI", iri: value };
+  }
+  if (value.startsWith("_:")) {
+    return { type: "BLANK", id: value.slice(2) };
+  }
+  return { type: "LITERAL", value };
+}