From 843234de44d0c03e5c19f2f0cd90bba93300c790 Mon Sep 17 00:00:00 2001
From: Ian Arawjo <fatso784@gmail.com>
Date: Wed, 3 May 2023 20:14:37 -0400
Subject: [PATCH] Query LLMs simultaneously

---
 python-backend/app.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/python-backend/app.py b/python-backend/app.py
index 62b7d2b..0d2027b 100644
--- a/python-backend/app.py
+++ b/python-backend/app.py
@@ -1,4 +1,4 @@
-import json, os
+import json, os, asyncio
 from dataclasses import dataclass
 from statistics import mean, median, stdev
 from flask import Flask, request, jsonify
@@ -187,8 +187,7 @@ async def queryLLM():
     responses = {}
     params = data['params'] if 'params' in data else {}
 
-    for llm in llms:
-
+    async def query(llm: str) -> list:
         # Check that storage path is valid:
         cache_filepath = os.path.join('cache', f"{data['id']}-{str(llm.name)}.json")
         if not is_valid_filepath(cache_filepath):
@@ -199,14 +198,28 @@ async def queryLLM():
 
         # Prompt the LLM with all permutations of the input prompt template:
         # NOTE: If the responses are already cache'd, this just loads them (no LLM is queried, saving $$$)
-        responses[llm] = []
+        resps = []
         try:
+            print(f'Querying {llm}...')
             async for response in prompter.gen_responses(properties=data['vars'], llm=llm, **params):
-                responses[llm].append(response)
+                resps.append(response)
         except Exception as e:
             print('error generating responses:', e)
             raise e
-            return jsonify({'error': str(e)})
+        
+        return {'llm': llm, 'responses': resps}
+            
+    try:
+        # Request responses simultaneously across LLMs
+        tasks = [query(llm) for llm in llms]
+
+        # Await the responses from all queried LLMs
+        llm_results = await asyncio.gather(*tasks)
+        for item in llm_results:
+            responses[item['llm']] = item['responses']
+
+    except Exception as e:
+        return jsonify({'error': str(e)})
 
     # Convert the responses into a more standardized format with less information
     res = [