Merge branch 'main' of https://github.com/biothings/bte_trapi_query_g…

…raph_handler into dev
biothings · Aug 6, 2024 · 3afe25f · 3afe25f
2 parents 1a73747 + f6eaa70
commit 3afe25f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/src/graph/kg_node.ts b/src/graph/kg_node.ts
@@ -8,12 +8,14 @@ export interface KGNodeInfo {
   curies: string[];
   primaryCurie: string;
   qNodeID: string;
+  originalCurie?: string;
 }
 
 export default class KGNode {
   id: string;
   primaryCurie: string;
   qNodeID: string;
+  originalCurie: string;
   curies: string[];
   names: string[];
   semanticType: string[];
@@ -36,6 +38,9 @@ export default class KGNode {
     this.targetNodes = new Set();
     this.sourceQNodeIDs = new Set();
     this.targetQNodeIDs = new Set();
+
+    // store original curie to output `query_id bte#815`
+    this.originalCurie = info.originalCurie;
   }
 
   addSourceNode(kgNodeID: string): void {

diff --git a/src/index.ts b/src/index.ts
@@ -319,6 +319,18 @@ export default class TRAPIQueryHandler {
     this.finalizedResults = fixedResults;
   }
 
+  appendOriginalCuriesToResults(results: TrapiResult[]): void {
+    results.forEach(result => {
+      Object.entries(result.node_bindings).forEach(([_, bindings]) => {
+        bindings.forEach(binding => {
+          if (this.bteGraph.nodes[binding.id].originalCurie && this.bteGraph.nodes[binding.id].originalCurie !== binding.id) {
+            binding.query_id = this.bteGraph.nodes[binding.id].originalCurie;
+          }
+        })
+      })
+    })
+  }
+
   async addQueryNodes(): Promise<void> {
     const qNodeIDsByOriginalID: Map<string, TrapiQNode> = new Map();
     const curiesToResolve = [
@@ -340,6 +352,7 @@ export default class TRAPIQueryHandler {
         this.bteGraph.nodes[resolvedEntity.primaryID] = new KGNode(resolvedEntity.primaryID, {
           primaryCurie: resolvedEntity.primaryID,
           qNodeID: qNodeIDsByOriginalID[originalCurie],
+          originalCurie: originalCurie,
           curies: resolvedEntity.equivalentIDs,
           names: resolvedEntity.labelAliases,
           semanticType: category ? [category] : ['biolink:NamedThing'],
@@ -727,6 +740,8 @@ export default class TRAPIQueryHandler {
     this.createSubclassSupportGraphs();
     // prune bteGraph
     this.bteGraph.prune(this.finalizedResults, this.auxGraphs);
+    // add original curies to results
+    this.appendOriginalCuriesToResults(this.finalizedResults);
     this.bteGraph.notify();
 
     // Attempt to enrich results with PFOCR figures