From a655ea15bde4cdc6c93d9369e593faa4103debf8 Mon Sep 17 00:00:00 2001
From: Rohan Juneja <rohanj2006@gmail.com>
Date: Fri, 26 Apr 2024 17:16:22 -0700
Subject: [PATCH] configurable setting for query duration on templates

---
 src/inferred_mode/inferred_mode.ts   | 13 +++++++------
 src/inferred_mode/template_lookup.ts |  5 ++++-
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/inferred_mode/inferred_mode.ts b/src/inferred_mode/inferred_mode.ts
index b98f92fa..86b651ed 100644
--- a/src/inferred_mode/inferred_mode.ts
+++ b/src/inferred_mode/inferred_mode.ts
@@ -214,7 +214,7 @@ export default class InferredQueryHandler {
   async createQueries(qEdge: TrapiQEdge, qSubject: TrapiQNode, qObject: TrapiQNode): Promise<FilledTemplate[]> {
     const templates = await this.findTemplates(qEdge, qSubject, qObject);
     // combine creative query with templates
-    const subQueries = templates.map(({ template, queryGraph }) => {
+    const subQueries = templates.map(({ template, queryGraph, durationMin }) => {
       queryGraph.nodes.creativeQuerySubject.categories = [
         ...new Set([...queryGraph.nodes.creativeQuerySubject.categories, ...qSubject.categories]),
       ];
@@ -244,7 +244,7 @@ export default class InferredQueryHandler {
         delete queryGraph.nodes.creativeQueryObject.ids;
       }
 
-      return { template, queryGraph };
+      return { template, queryGraph, durationMin };
     });
 
     return subQueries;
@@ -525,12 +525,13 @@ export default class InferredQueryHandler {
 
     
     const MAX_TIME = 4.5 * 60 * 1000; // 4 minutes
-    const QUERY_TIME = 2.5 * 60 * 1000; // 2.5 minutes
+    const DEFAULT_QUERY_TIME = 2.5 * 60 * 1000; // 2.5 minutes
     const start = Date.now();
 
-    await async.eachOfSeries(subQueries, async ({ template, queryGraph }, i) => {
-      if (Date.now() - start > MAX_TIME - QUERY_TIME) {
-        debug(`Skipping template because the query has been running for ${(Date.now() - start) / 1000} seconds, and this template is projected to take ${QUERY_TIME / 1000} seconds`);
+    await async.eachOfSeries(subQueries, async ({ template, queryGraph, durationMin }, i) => {
+      const queryTime = durationMin * 60 * 1000 ?? DEFAULT_QUERY_TIME;
+      if (Date.now() - start > MAX_TIME - queryTime) {
+        debug(`Skipping template because the query has been running for ${(Date.now() - start) / 1000} seconds, and this template is projected to take ${queryTime / 1000} seconds`);
         return;
       }
       const span = Telemetry.startSpan({ description: 'creativeTemplate' });
diff --git a/src/inferred_mode/template_lookup.ts b/src/inferred_mode/template_lookup.ts
index e32973cb..658b5897 100644
--- a/src/inferred_mode/template_lookup.ts
+++ b/src/inferred_mode/template_lookup.ts
@@ -17,6 +17,7 @@ export interface TemplateLookup {
 export interface MatchedTemplate {
   template: string;
   queryGraph: TrapiQueryGraph;
+  durationMin?: number;
 }
 
 export interface TemplateGroup {
@@ -77,9 +78,11 @@ export async function getTemplates(lookups: TemplateLookup[]): Promise<MatchedTe
     return matches;
   }, [] as string[]);
   return await async.map(matchingTemplatePaths, async (templatePath: string) => {
+    const templateData = JSON.parse(await fs.readFile(templatePath, { encoding: 'utf8' }));
     return {
       template: templatePath.substring(templatePath.lastIndexOf('/') + 1),
-      queryGraph: JSON.parse(await fs.readFile(templatePath, { encoding: 'utf8' })).message.query_graph,
+      queryGraph: templateData.message.query_graph,
+      durationMin: templateData.durationMin
     };
   });
 }